
<!DOCTYPE html>

<html lang="en">
  <head>
    <meta charset="utf-8" />
    <meta name="viewport" content="width=device-width, initial-scale=1.0" />
    <title>tigramite.plotting &#8212; Tigramite 5.2 documentation</title>
    <link rel="stylesheet" type="text/css" href="../../_static/pygments.css" />
    <link rel="stylesheet" type="text/css" href="../../_static/alabaster.css" />
    <script data-url_root="../../" id="documentation_options" src="../../_static/documentation_options.js"></script>
    <script src="../../_static/jquery.js"></script>
    <script src="../../_static/underscore.js"></script>
    <script src="../../_static/_sphinx_javascript_frameworks_compat.js"></script>
    <script src="../../_static/doctools.js"></script>
    <link rel="index" title="Index" href="../../genindex.html" />
    <link rel="search" title="Search" href="../../search.html" />
   
  <link rel="stylesheet" href="../../_static/custom.css" type="text/css" />
  
  
  <meta name="viewport" content="width=device-width, initial-scale=0.9, maximum-scale=0.9" />

  </head><body>
  

    <div class="document">
      <div class="documentwrapper">
        <div class="bodywrapper">
          

          <div class="body" role="main">
            
  <h1>Source code for tigramite.plotting</h1><div class="highlight"><pre>
<span></span><span class="sd">&quot;&quot;&quot;Tigramite plotting package.&quot;&quot;&quot;</span>

<span class="c1"># Author: Jakob Runge &lt;jakob@jakob-runge.com&gt;</span>
<span class="c1">#</span>
<span class="c1"># License: GNU General Public License v3.0</span>

<span class="kn">import</span> <span class="nn">numpy</span> <span class="k">as</span> <span class="nn">np</span>
<span class="kn">import</span> <span class="nn">json</span><span class="o">,</span> <span class="nn">warnings</span><span class="o">,</span> <span class="nn">os</span><span class="o">,</span> <span class="nn">pathlib</span>
<span class="kn">import</span> <span class="nn">matplotlib</span>
<span class="kn">import</span> <span class="nn">networkx</span> <span class="k">as</span> <span class="nn">nx</span>
<span class="kn">from</span> <span class="nn">matplotlib.colors</span> <span class="kn">import</span> <span class="n">ListedColormap</span>
<span class="kn">import</span> <span class="nn">matplotlib.transforms</span> <span class="k">as</span> <span class="nn">transforms</span>
<span class="kn">from</span> <span class="nn">matplotlib</span> <span class="kn">import</span> <span class="n">pyplot</span><span class="p">,</span> <span class="n">ticker</span>
<span class="kn">from</span> <span class="nn">matplotlib.ticker</span> <span class="kn">import</span> <span class="n">FormatStrFormatter</span>
<span class="kn">import</span> <span class="nn">matplotlib.patches</span> <span class="k">as</span> <span class="nn">mpatches</span>
<span class="kn">from</span> <span class="nn">matplotlib.collections</span> <span class="kn">import</span> <span class="n">PatchCollection</span>
<span class="kn">from</span> <span class="nn">mpl_toolkits.axes_grid1</span> <span class="kn">import</span> <span class="n">make_axes_locatable</span>
<span class="kn">import</span> <span class="nn">sys</span>
<span class="kn">from</span> <span class="nn">operator</span> <span class="kn">import</span> <span class="n">sub</span>
<span class="kn">import</span> <span class="nn">tigramite.data_processing</span> <span class="k">as</span> <span class="nn">pp</span>
<span class="kn">from</span> <span class="nn">copy</span> <span class="kn">import</span> <span class="n">deepcopy</span>
<span class="kn">import</span> <span class="nn">matplotlib.path</span> <span class="k">as</span> <span class="nn">mpath</span>
<span class="kn">import</span> <span class="nn">matplotlib.patheffects</span> <span class="k">as</span> <span class="nn">PathEffects</span>
<span class="kn">from</span> <span class="nn">mpl_toolkits.axisartist.axislines</span> <span class="kn">import</span> <span class="n">Axes</span>
<span class="kn">import</span> <span class="nn">csv</span>


<span class="c1"># TODO: Add proper docstrings to internal functions...</span>


<span class="k">def</span> <span class="nf">_par_corr_trafo</span><span class="p">(</span><span class="n">cmi</span><span class="p">):</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;Transformation of CMI to partial correlation scale.&quot;&quot;&quot;</span>

    <span class="c1"># Set negative values to small positive number</span>
    <span class="c1"># (zero would be interpreted as non-significant in some functions)</span>
    <span class="k">if</span> <span class="n">np</span><span class="o">.</span><span class="n">ndim</span><span class="p">(</span><span class="n">cmi</span><span class="p">)</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
        <span class="k">if</span> <span class="n">cmi</span> <span class="o">&lt;</span> <span class="mf">0.0</span><span class="p">:</span>
            <span class="n">cmi</span> <span class="o">=</span> <span class="mf">1e-8</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">cmi</span><span class="p">[</span><span class="n">cmi</span> <span class="o">&lt;</span> <span class="mf">0.0</span><span class="p">]</span> <span class="o">=</span> <span class="mf">1e-8</span>

    <span class="k">return</span> <span class="n">np</span><span class="o">.</span><span class="n">sqrt</span><span class="p">(</span><span class="mf">1.0</span> <span class="o">-</span> <span class="n">np</span><span class="o">.</span><span class="n">exp</span><span class="p">(</span><span class="o">-</span><span class="mf">2.0</span> <span class="o">*</span> <span class="n">cmi</span><span class="p">))</span>


<span class="k">def</span> <span class="nf">_par_corr_to_cmi</span><span class="p">(</span><span class="n">par_corr</span><span class="p">):</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;Transformation of partial correlation to CMI scale.&quot;&quot;&quot;</span>

    <span class="k">return</span> <span class="o">-</span><span class="mf">0.5</span> <span class="o">*</span> <span class="n">np</span><span class="o">.</span><span class="n">log</span><span class="p">(</span><span class="mf">1.0</span> <span class="o">-</span> <span class="n">par_corr</span> <span class="o">**</span> <span class="mi">2</span><span class="p">)</span>


<span class="k">def</span> <span class="nf">_myround</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">base</span><span class="o">=</span><span class="mi">5</span><span class="p">,</span> <span class="n">round_mode</span><span class="o">=</span><span class="s2">&quot;updown&quot;</span><span class="p">):</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;Rounds x to a float with precision base.&quot;&quot;&quot;</span>

    <span class="k">if</span> <span class="n">round_mode</span> <span class="o">==</span> <span class="s2">&quot;updown&quot;</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">base</span> <span class="o">*</span> <span class="nb">round</span><span class="p">(</span><span class="nb">float</span><span class="p">(</span><span class="n">x</span><span class="p">)</span> <span class="o">/</span> <span class="n">base</span><span class="p">)</span>
    <span class="k">elif</span> <span class="n">round_mode</span> <span class="o">==</span> <span class="s2">&quot;down&quot;</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">base</span> <span class="o">*</span> <span class="n">np</span><span class="o">.</span><span class="n">floor</span><span class="p">(</span><span class="nb">float</span><span class="p">(</span><span class="n">x</span><span class="p">)</span> <span class="o">/</span> <span class="n">base</span><span class="p">)</span>
    <span class="k">elif</span> <span class="n">round_mode</span> <span class="o">==</span> <span class="s2">&quot;up&quot;</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">base</span> <span class="o">*</span> <span class="n">np</span><span class="o">.</span><span class="n">ceil</span><span class="p">(</span><span class="nb">float</span><span class="p">(</span><span class="n">x</span><span class="p">)</span> <span class="o">/</span> <span class="n">base</span><span class="p">)</span>

    <span class="k">return</span> <span class="n">base</span> <span class="o">*</span> <span class="nb">round</span><span class="p">(</span><span class="nb">float</span><span class="p">(</span><span class="n">x</span><span class="p">)</span> <span class="o">/</span> <span class="n">base</span><span class="p">)</span>


<span class="k">def</span> <span class="nf">_make_nice_axes</span><span class="p">(</span><span class="n">ax</span><span class="p">,</span> <span class="n">where</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">skip</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span> <span class="n">color</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;Makes nice axes.&quot;&quot;&quot;</span>

    <span class="k">if</span> <span class="n">where</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">where</span> <span class="o">=</span> <span class="p">[</span><span class="s2">&quot;left&quot;</span><span class="p">,</span> <span class="s2">&quot;bottom&quot;</span><span class="p">]</span>
    <span class="k">if</span> <span class="n">color</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">color</span> <span class="o">=</span> <span class="p">{</span><span class="s2">&quot;left&quot;</span><span class="p">:</span> <span class="s2">&quot;black&quot;</span><span class="p">,</span> <span class="s2">&quot;right&quot;</span><span class="p">:</span> <span class="s2">&quot;black&quot;</span><span class="p">,</span> <span class="s2">&quot;bottom&quot;</span><span class="p">:</span> <span class="s2">&quot;black&quot;</span><span class="p">,</span> <span class="s2">&quot;top&quot;</span><span class="p">:</span> <span class="s2">&quot;black&quot;</span><span class="p">}</span>

    <span class="k">if</span> <span class="nb">type</span><span class="p">(</span><span class="n">skip</span><span class="p">)</span> <span class="o">==</span> <span class="nb">int</span><span class="p">:</span>
        <span class="n">skip_x</span> <span class="o">=</span> <span class="n">skip_y</span> <span class="o">=</span> <span class="n">skip</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">skip_x</span> <span class="o">=</span> <span class="n">skip</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
        <span class="n">skip_y</span> <span class="o">=</span> <span class="n">skip</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span>

    <span class="k">for</span> <span class="n">loc</span><span class="p">,</span> <span class="n">spine</span> <span class="ow">in</span> <span class="n">ax</span><span class="o">.</span><span class="n">spines</span><span class="o">.</span><span class="n">items</span><span class="p">():</span>
        <span class="k">if</span> <span class="n">loc</span> <span class="ow">in</span> <span class="n">where</span><span class="p">:</span>
            <span class="n">spine</span><span class="o">.</span><span class="n">set_position</span><span class="p">((</span><span class="s2">&quot;outward&quot;</span><span class="p">,</span> <span class="mi">5</span><span class="p">))</span>  <span class="c1"># outward by 10 points</span>
            <span class="n">spine</span><span class="o">.</span><span class="n">set_color</span><span class="p">(</span><span class="n">color</span><span class="p">[</span><span class="n">loc</span><span class="p">])</span>
            <span class="k">if</span> <span class="n">loc</span> <span class="o">==</span> <span class="s2">&quot;left&quot;</span> <span class="ow">or</span> <span class="n">loc</span> <span class="o">==</span> <span class="s2">&quot;right&quot;</span><span class="p">:</span>
                <span class="n">pyplot</span><span class="o">.</span><span class="n">setp</span><span class="p">(</span><span class="n">ax</span><span class="o">.</span><span class="n">get_yticklines</span><span class="p">(),</span> <span class="n">color</span><span class="o">=</span><span class="n">color</span><span class="p">[</span><span class="n">loc</span><span class="p">])</span>
                <span class="n">pyplot</span><span class="o">.</span><span class="n">setp</span><span class="p">(</span><span class="n">ax</span><span class="o">.</span><span class="n">get_yticklabels</span><span class="p">(),</span> <span class="n">color</span><span class="o">=</span><span class="n">color</span><span class="p">[</span><span class="n">loc</span><span class="p">])</span>
            <span class="k">if</span> <span class="n">loc</span> <span class="o">==</span> <span class="s2">&quot;top&quot;</span> <span class="ow">or</span> <span class="n">loc</span> <span class="o">==</span> <span class="s2">&quot;bottom&quot;</span><span class="p">:</span>
                <span class="n">pyplot</span><span class="o">.</span><span class="n">setp</span><span class="p">(</span><span class="n">ax</span><span class="o">.</span><span class="n">get_xticklines</span><span class="p">(),</span> <span class="n">color</span><span class="o">=</span><span class="n">color</span><span class="p">[</span><span class="n">loc</span><span class="p">])</span>
        <span class="k">elif</span> <span class="n">loc</span> <span class="ow">in</span> <span class="p">[</span>
            <span class="n">item</span> <span class="k">for</span> <span class="n">item</span> <span class="ow">in</span> <span class="p">[</span><span class="s2">&quot;left&quot;</span><span class="p">,</span> <span class="s2">&quot;bottom&quot;</span><span class="p">,</span> <span class="s2">&quot;right&quot;</span><span class="p">,</span> <span class="s2">&quot;top&quot;</span><span class="p">]</span> <span class="k">if</span> <span class="n">item</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">where</span>
        <span class="p">]:</span>
            <span class="n">spine</span><span class="o">.</span><span class="n">set_color</span><span class="p">(</span><span class="s2">&quot;none&quot;</span><span class="p">)</span>  <span class="c1"># don&#39;t draw spine</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;unknown spine location: </span><span class="si">%s</span><span class="s2">&quot;</span> <span class="o">%</span> <span class="n">loc</span><span class="p">)</span>

    <span class="c1"># ax.xaxis.get_major_formatter().set_useOffset(False)</span>

    <span class="c1"># turn off ticks where there is no spine</span>
    <span class="k">if</span> <span class="s2">&quot;top&quot;</span> <span class="ow">in</span> <span class="n">where</span> <span class="ow">and</span> <span class="s2">&quot;bottom&quot;</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">where</span><span class="p">:</span>
        <span class="n">ax</span><span class="o">.</span><span class="n">xaxis</span><span class="o">.</span><span class="n">set_ticks_position</span><span class="p">(</span><span class="s2">&quot;top&quot;</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">skip_x</span> <span class="o">&gt;</span> <span class="mi">1</span><span class="p">:</span>
            <span class="n">ax</span><span class="o">.</span><span class="n">set_xticks</span><span class="p">(</span><span class="n">ax</span><span class="o">.</span><span class="n">get_xticks</span><span class="p">()[::</span><span class="n">skip_x</span><span class="p">])</span>
    <span class="k">elif</span> <span class="s2">&quot;bottom&quot;</span> <span class="ow">in</span> <span class="n">where</span><span class="p">:</span>
        <span class="n">ax</span><span class="o">.</span><span class="n">xaxis</span><span class="o">.</span><span class="n">set_ticks_position</span><span class="p">(</span><span class="s2">&quot;bottom&quot;</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">skip_x</span> <span class="o">&gt;</span> <span class="mi">1</span><span class="p">:</span>
            <span class="n">ax</span><span class="o">.</span><span class="n">set_xticks</span><span class="p">(</span><span class="n">ax</span><span class="o">.</span><span class="n">get_xticks</span><span class="p">()[::</span><span class="n">skip_x</span><span class="p">])</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">ax</span><span class="o">.</span><span class="n">xaxis</span><span class="o">.</span><span class="n">set_ticks_position</span><span class="p">(</span><span class="s2">&quot;none&quot;</span><span class="p">)</span>
        <span class="n">ax</span><span class="o">.</span><span class="n">xaxis</span><span class="o">.</span><span class="n">set_ticklabels</span><span class="p">([])</span>
    <span class="k">if</span> <span class="s2">&quot;right&quot;</span> <span class="ow">in</span> <span class="n">where</span> <span class="ow">and</span> <span class="s2">&quot;left&quot;</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">where</span><span class="p">:</span>
        <span class="n">ax</span><span class="o">.</span><span class="n">yaxis</span><span class="o">.</span><span class="n">set_ticks_position</span><span class="p">(</span><span class="s2">&quot;right&quot;</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">skip_y</span> <span class="o">&gt;</span> <span class="mi">1</span><span class="p">:</span>
            <span class="n">ax</span><span class="o">.</span><span class="n">set_yticks</span><span class="p">(</span><span class="n">ax</span><span class="o">.</span><span class="n">get_yticks</span><span class="p">()[::</span><span class="n">skip_y</span><span class="p">])</span>
    <span class="k">elif</span> <span class="s2">&quot;left&quot;</span> <span class="ow">in</span> <span class="n">where</span><span class="p">:</span>
        <span class="n">ax</span><span class="o">.</span><span class="n">yaxis</span><span class="o">.</span><span class="n">set_ticks_position</span><span class="p">(</span><span class="s2">&quot;left&quot;</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">skip_y</span> <span class="o">&gt;</span> <span class="mi">1</span><span class="p">:</span>
            <span class="n">ax</span><span class="o">.</span><span class="n">set_yticks</span><span class="p">(</span><span class="n">ax</span><span class="o">.</span><span class="n">get_yticks</span><span class="p">()[::</span><span class="n">skip_y</span><span class="p">])</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">ax</span><span class="o">.</span><span class="n">yaxis</span><span class="o">.</span><span class="n">set_ticks_position</span><span class="p">(</span><span class="s2">&quot;none&quot;</span><span class="p">)</span>
        <span class="n">ax</span><span class="o">.</span><span class="n">yaxis</span><span class="o">.</span><span class="n">set_ticklabels</span><span class="p">([])</span>

    <span class="n">ax</span><span class="o">.</span><span class="n">patch</span><span class="o">.</span><span class="n">set_alpha</span><span class="p">(</span><span class="mf">0.0</span><span class="p">)</span>


<span class="k">def</span> <span class="nf">_get_absmax</span><span class="p">(</span><span class="n">val_matrix</span><span class="p">):</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;Get value at absolute maximum in lag function array.</span>
<span class="sd">    For an (N, N, tau)-array this comutes the lag of the absolute maximum</span>
<span class="sd">    along the tau-axis and stores the (positive or negative) value in</span>
<span class="sd">    the (N,N)-array absmax.&quot;&quot;&quot;</span>

    <span class="n">absmax_indices</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">abs</span><span class="p">(</span><span class="n">val_matrix</span><span class="p">)</span><span class="o">.</span><span class="n">argmax</span><span class="p">(</span><span class="n">axis</span><span class="o">=</span><span class="mi">2</span><span class="p">)</span>
    <span class="n">i</span><span class="p">,</span> <span class="n">j</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">indices</span><span class="p">(</span><span class="n">val_matrix</span><span class="o">.</span><span class="n">shape</span><span class="p">[:</span><span class="mi">2</span><span class="p">])</span>

    <span class="k">return</span> <span class="n">val_matrix</span><span class="p">[</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">,</span> <span class="n">absmax_indices</span><span class="p">]</span>


<span class="k">def</span> <span class="nf">_add_timeseries</span><span class="p">(</span>
        <span class="n">dataframe</span><span class="p">,</span>
        <span class="n">fig_axes</span><span class="p">,</span>
        <span class="n">grey_masked_samples</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
        <span class="n">show_meanline</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
        <span class="n">data_linewidth</span><span class="o">=</span><span class="mf">1.0</span><span class="p">,</span>
        <span class="n">color</span><span class="o">=</span><span class="s2">&quot;black&quot;</span><span class="p">,</span>
        <span class="n">alpha</span><span class="o">=</span><span class="mf">1.</span><span class="p">,</span>
        <span class="n">grey_alpha</span><span class="o">=</span><span class="mf">1.0</span><span class="p">,</span>
        <span class="n">selected_dataset</span><span class="o">=</span><span class="mi">0</span><span class="p">,</span>
        <span class="n">selected_variables</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
<span class="p">):</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;Adds a time series plot to an axis.</span>
<span class="sd">    Plot of dataseries is added to axis. Allows for proper visualization of</span>
<span class="sd">    masked data.</span>

<span class="sd">    Parameters</span>
<span class="sd">    ----------</span>
<span class="sd">    fig : figure instance</span>
<span class="sd">        Figure instance.</span>
<span class="sd">    axes : axis instance</span>
<span class="sd">        Either gridded axis object or single axis instance.</span>
<span class="sd">    grey_masked_samples : bool, optional (default: False)</span>
<span class="sd">        Whether to mark masked samples by grey fills (&#39;fill&#39;) or grey data</span>
<span class="sd">        (&#39;data&#39;).</span>
<span class="sd">    show_meanline : bool</span>
<span class="sd">        Show mean of data as horizontal line.</span>
<span class="sd">    data_linewidth : float, optional (default: 1.)</span>
<span class="sd">        Linewidth.</span>
<span class="sd">    color : str, optional (default: black)</span>
<span class="sd">        Line color.</span>
<span class="sd">    alpha : float</span>
<span class="sd">        Alpha opacity.</span>
<span class="sd">    grey_alpha : float, optional (default: 1.)</span>
<span class="sd">        Opacity of fill_between.</span>
<span class="sd">    selected_dataset : int, optional (default: 0)</span>
<span class="sd">        In case of multiple datasets in dataframe, plot this one.</span>
<span class="sd">    selected_variables : list, optional (default: None)</span>
<span class="sd">        List of variables which to plot.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">fig</span><span class="p">,</span> <span class="n">axes</span> <span class="o">=</span> <span class="n">fig_axes</span>

    <span class="c1"># Read in all attributes from dataframe</span>
    <span class="n">data</span> <span class="o">=</span> <span class="n">dataframe</span><span class="o">.</span><span class="n">values</span><span class="p">[</span><span class="n">selected_dataset</span><span class="p">]</span>
    <span class="k">if</span> <span class="n">dataframe</span><span class="o">.</span><span class="n">mask</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">mask</span> <span class="o">=</span> <span class="n">dataframe</span><span class="o">.</span><span class="n">mask</span><span class="p">[</span><span class="n">selected_dataset</span><span class="p">]</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">mask</span> <span class="o">=</span> <span class="kc">None</span>

    <span class="n">missing_flag</span> <span class="o">=</span> <span class="n">dataframe</span><span class="o">.</span><span class="n">missing_flag</span>
    <span class="n">time</span> <span class="o">=</span> <span class="n">dataframe</span><span class="o">.</span><span class="n">datatime</span><span class="p">[</span><span class="n">selected_dataset</span><span class="p">]</span>
    <span class="n">T</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="n">time</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">selected_variables</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">selected_variables</span> <span class="o">=</span> <span class="nb">list</span><span class="p">(</span><span class="nb">range</span><span class="p">(</span><span class="n">dataframe</span><span class="o">.</span><span class="n">N</span><span class="p">))</span>

    <span class="n">nb_components</span> <span class="o">=</span> <span class="nb">sum</span><span class="p">([</span><span class="nb">len</span><span class="p">(</span><span class="n">dataframe</span><span class="o">.</span><span class="n">vector_vars</span><span class="p">[</span><span class="n">var</span><span class="p">])</span> <span class="k">for</span> <span class="n">var</span> <span class="ow">in</span> <span class="n">selected_variables</span><span class="p">])</span>

    <span class="k">for</span> <span class="n">j</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">nb_components</span><span class="p">):</span>

        <span class="n">ax</span> <span class="o">=</span> <span class="n">axes</span><span class="p">[</span><span class="n">j</span><span class="p">]</span>
        <span class="n">dataseries</span> <span class="o">=</span> <span class="n">data</span><span class="p">[:,</span> <span class="n">j</span><span class="p">]</span>

        <span class="k">if</span> <span class="n">missing_flag</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">dataseries_nomissing</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">ma</span><span class="o">.</span><span class="n">masked_where</span><span class="p">(</span>
                <span class="n">dataseries</span> <span class="o">==</span> <span class="n">missing_flag</span><span class="p">,</span> <span class="n">dataseries</span>
            <span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">dataseries_nomissing</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">ma</span><span class="o">.</span><span class="n">masked_where</span><span class="p">(</span>
                <span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="n">dataseries</span><span class="o">.</span><span class="n">shape</span><span class="p">),</span> <span class="n">dataseries</span>
            <span class="p">)</span>


        <span class="k">if</span> <span class="n">mask</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">maskseries</span> <span class="o">=</span> <span class="n">mask</span><span class="p">[:,</span> <span class="n">j</span><span class="p">]</span>

            <span class="n">maskdata</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">ma</span><span class="o">.</span><span class="n">masked_where</span><span class="p">(</span><span class="n">maskseries</span><span class="p">,</span> <span class="n">dataseries_nomissing</span><span class="p">)</span>

            <span class="k">if</span> <span class="n">grey_masked_samples</span> <span class="o">==</span> <span class="s2">&quot;fill&quot;</span><span class="p">:</span>
                <span class="n">ax</span><span class="o">.</span><span class="n">fill_between</span><span class="p">(</span>
                    <span class="n">time</span><span class="p">,</span>
                    <span class="n">maskdata</span><span class="o">.</span><span class="n">min</span><span class="p">(),</span>
                    <span class="n">maskdata</span><span class="o">.</span><span class="n">max</span><span class="p">(),</span>
                    <span class="n">where</span><span class="o">=</span><span class="n">maskseries</span><span class="p">,</span>
                    <span class="n">color</span><span class="o">=</span><span class="s2">&quot;grey&quot;</span><span class="p">,</span>
                    <span class="n">interpolate</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
                    <span class="n">linewidth</span><span class="o">=</span><span class="mf">0.0</span><span class="p">,</span>
                    <span class="n">alpha</span><span class="o">=</span><span class="n">grey_alpha</span><span class="p">,</span>
                <span class="p">)</span>
            <span class="k">elif</span> <span class="n">grey_masked_samples</span> <span class="o">==</span> <span class="s2">&quot;data&quot;</span><span class="p">:</span>
                <span class="n">ax</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span>
                    <span class="n">time</span><span class="p">,</span>
                    <span class="n">dataseries_nomissing</span><span class="p">,</span>
                    <span class="n">color</span><span class="o">=</span><span class="s2">&quot;grey&quot;</span><span class="p">,</span>
                    <span class="n">marker</span><span class="o">=</span><span class="s2">&quot;.&quot;</span><span class="p">,</span>
                    <span class="n">markersize</span><span class="o">=</span><span class="n">data_linewidth</span><span class="p">,</span>
                    <span class="n">linewidth</span><span class="o">=</span><span class="n">data_linewidth</span><span class="p">,</span>
                    <span class="n">clip_on</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
                    <span class="n">alpha</span><span class="o">=</span><span class="n">grey_alpha</span><span class="p">,</span>
                <span class="p">)</span>
            <span class="k">if</span> <span class="n">show_meanline</span><span class="p">:</span>
                <span class="n">ax</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span><span class="n">time</span><span class="p">,</span> <span class="n">maskdata</span><span class="o">.</span><span class="n">mean</span><span class="p">()</span> <span class="o">*</span> <span class="n">np</span><span class="o">.</span><span class="n">ones</span><span class="p">(</span><span class="n">T</span><span class="p">),</span> <span class="n">lw</span><span class="o">=</span><span class="n">data_linewidth</span> <span class="o">/</span> <span class="mf">2.</span><span class="p">,</span> <span class="n">color</span><span class="o">=</span><span class="n">color</span><span class="p">)</span>

            <span class="n">ax</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span>
                <span class="n">time</span><span class="p">,</span>
                <span class="n">maskdata</span><span class="p">,</span>
                <span class="n">color</span><span class="o">=</span><span class="n">color</span><span class="p">,</span>
                <span class="n">linewidth</span><span class="o">=</span><span class="n">data_linewidth</span><span class="p">,</span>
                <span class="n">marker</span><span class="o">=</span><span class="s2">&quot;.&quot;</span><span class="p">,</span>
                <span class="n">markersize</span><span class="o">=</span><span class="n">data_linewidth</span><span class="p">,</span>
                <span class="n">clip_on</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
                <span class="n">alpha</span><span class="o">=</span><span class="n">alpha</span><span class="p">,</span>
            <span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="k">if</span> <span class="n">show_meanline</span><span class="p">:</span>
                <span class="n">ax</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span><span class="n">time</span><span class="p">,</span> <span class="n">dataseries_nomissing</span><span class="o">.</span><span class="n">mean</span><span class="p">()</span> <span class="o">*</span> <span class="n">np</span><span class="o">.</span><span class="n">ones</span><span class="p">(</span><span class="n">T</span><span class="p">),</span> <span class="n">lw</span><span class="o">=</span><span class="n">data_linewidth</span> <span class="o">/</span> <span class="mf">2.</span><span class="p">,</span> <span class="n">color</span><span class="o">=</span><span class="n">color</span><span class="p">)</span>

            <span class="n">ax</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span>
                <span class="n">time</span><span class="p">,</span>
                <span class="n">dataseries_nomissing</span><span class="p">,</span>
                <span class="n">color</span><span class="o">=</span><span class="n">color</span><span class="p">,</span>
                <span class="n">linewidth</span><span class="o">=</span><span class="n">data_linewidth</span><span class="p">,</span>
                <span class="n">clip_on</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
                <span class="n">alpha</span><span class="o">=</span><span class="n">alpha</span><span class="p">,</span>
            <span class="p">)</span>


<div class="viewcode-block" id="plot_timeseries"><a class="viewcode-back" href="../../index.html#tigramite.plotting.plot_timeseries">[docs]</a><span class="k">def</span> <span class="nf">plot_timeseries</span><span class="p">(</span>
        <span class="n">dataframe</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
        <span class="n">save_name</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
        <span class="n">fig_axes</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
        <span class="n">figsize</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
        <span class="n">var_units</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
        <span class="n">time_label</span><span class="o">=</span><span class="s2">&quot;&quot;</span><span class="p">,</span>
        <span class="n">grey_masked_samples</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
        <span class="n">show_meanline</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
        <span class="n">data_linewidth</span><span class="o">=</span><span class="mf">1.0</span><span class="p">,</span>
        <span class="n">skip_ticks_data_x</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span>
        <span class="n">skip_ticks_data_y</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span>
        <span class="n">label_fontsize</span><span class="o">=</span><span class="mi">10</span><span class="p">,</span>
        <span class="n">color</span><span class="o">=</span><span class="s1">&#39;black&#39;</span><span class="p">,</span>
        <span class="n">alpha</span><span class="o">=</span><span class="mf">1.</span><span class="p">,</span>
        <span class="n">tick_label_size</span><span class="o">=</span><span class="mi">6</span><span class="p">,</span>
        <span class="n">selected_dataset</span><span class="o">=</span><span class="mi">0</span><span class="p">,</span>
        <span class="n">adjust_plot</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
        <span class="n">selected_variables</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
<span class="p">):</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;Create and save figure of stacked panels with time series.</span>

<span class="sd">    Parameters</span>
<span class="sd">    ----------</span>
<span class="sd">    dataframe : data object, optional</span>
<span class="sd">        This is the Tigramite dataframe object. It has the attributes</span>
<span class="sd">        dataframe.values yielding a np array of shape (observations T,</span>
<span class="sd">        variables N) and optionally a mask of the same shape.</span>
<span class="sd">    save_name : str, optional (default: None)</span>
<span class="sd">        Name of figure file to save figure. If None, figure is shown in window.</span>
<span class="sd">    fig_axes : subplots instance, optional (default: None)</span>
<span class="sd">        Figure and axes instance. If None they are created as</span>
<span class="sd">        fig, axes = pyplot.subplots(N,...)</span>
<span class="sd">    figsize : tuple of floats, optional (default: None)</span>
<span class="sd">        Figure size if new figure is created. If None, default pyplot figsize</span>
<span class="sd">        is used.</span>
<span class="sd">    var_units : list of str, optional (default: None)</span>
<span class="sd">        Units of variables.</span>
<span class="sd">    time_label : str, optional (default: &#39;&#39;)</span>
<span class="sd">        Label of time axis.</span>
<span class="sd">    grey_masked_samples : bool, optional (default: False)</span>
<span class="sd">        Whether to mark masked samples by grey fills (&#39;fill&#39;) or grey data</span>
<span class="sd">        (&#39;data&#39;).</span>
<span class="sd">    show_meanline : bool, optional (default: False)</span>
<span class="sd">        Whether to plot a horizontal line at the mean.</span>
<span class="sd">    data_linewidth : float, optional (default: 1.)</span>
<span class="sd">        Linewidth.</span>
<span class="sd">    skip_ticks_data_x : int, optional (default: 1)</span>
<span class="sd">        Skip every other tickmark.</span>
<span class="sd">    skip_ticks_data_y : int, optional (default: 2)</span>
<span class="sd">        Skip every other tickmark.</span>
<span class="sd">    label_fontsize : int, optional (default: 10)</span>
<span class="sd">        Fontsize of variable labels.</span>
<span class="sd">    tick_label_size : int, optional (default: 6)</span>
<span class="sd">        Fontsize of tick labels.</span>
<span class="sd">    color : str, optional (default: black)</span>
<span class="sd">        Line color.</span>
<span class="sd">    alpha : float</span>
<span class="sd">        Alpha opacity.</span>
<span class="sd">    selected_dataset : int, optional (default: 0)</span>
<span class="sd">        In case of multiple datasets in dataframe, plot this one.</span>
<span class="sd">    selected_variables : list, optional (default: None)</span>
<span class="sd">        List of variables which to plot.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="n">var_names</span> <span class="o">=</span> <span class="n">dataframe</span><span class="o">.</span><span class="n">var_names</span>
    <span class="n">time</span> <span class="o">=</span> <span class="n">dataframe</span><span class="o">.</span><span class="n">datatime</span><span class="p">[</span><span class="n">selected_dataset</span><span class="p">]</span>

    <span class="n">N</span> <span class="o">=</span> <span class="n">dataframe</span><span class="o">.</span><span class="n">N</span>

    <span class="k">if</span> <span class="n">selected_variables</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">selected_variables</span> <span class="o">=</span> <span class="nb">list</span><span class="p">(</span><span class="nb">range</span><span class="p">(</span><span class="n">N</span><span class="p">))</span>

    <span class="n">nb_components_per_var</span> <span class="o">=</span> <span class="p">[</span><span class="nb">len</span><span class="p">(</span><span class="n">dataframe</span><span class="o">.</span><span class="n">vector_vars</span><span class="p">[</span><span class="n">var</span><span class="p">])</span> <span class="k">for</span> <span class="n">var</span> <span class="ow">in</span> <span class="n">selected_variables</span><span class="p">]</span>
    <span class="n">N_index</span> <span class="o">=</span> <span class="p">[</span><span class="nb">sum</span><span class="p">(</span><span class="n">nb_components_per_var</span><span class="p">[:</span><span class="n">i</span><span class="p">])</span> <span class="k">for</span> <span class="n">i</span><span class="p">,</span> <span class="n">el</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">nb_components_per_var</span><span class="p">)]</span>
    <span class="n">nb_components</span> <span class="o">=</span> <span class="nb">sum</span><span class="p">(</span><span class="n">nb_components_per_var</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">var_units</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">var_units</span> <span class="o">=</span> <span class="p">[</span><span class="s2">&quot;&quot;</span> <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">N</span><span class="p">)]</span>

    <span class="k">if</span> <span class="n">fig_axes</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">fig</span><span class="p">,</span> <span class="n">axes</span> <span class="o">=</span> <span class="n">pyplot</span><span class="o">.</span><span class="n">subplots</span><span class="p">(</span><span class="n">nb_components</span><span class="p">,</span> <span class="n">sharex</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span> <span class="n">figsize</span><span class="o">=</span><span class="n">figsize</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">fig</span><span class="p">,</span> <span class="n">axes</span> <span class="o">=</span> <span class="n">fig_axes</span>

    <span class="k">if</span> <span class="n">adjust_plot</span><span class="p">:</span>
        <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">nb_components</span><span class="p">):</span>

            <span class="n">ax</span> <span class="o">=</span> <span class="n">axes</span><span class="p">[</span><span class="n">i</span><span class="p">]</span>

            <span class="k">if</span> <span class="p">(</span><span class="n">i</span> <span class="o">==</span> <span class="n">nb_components</span> <span class="o">-</span> <span class="mi">1</span><span class="p">):</span>
                <span class="n">_make_nice_axes</span><span class="p">(</span>
                    <span class="n">ax</span><span class="p">,</span> <span class="n">where</span><span class="o">=</span><span class="p">[</span><span class="s2">&quot;left&quot;</span><span class="p">,</span> <span class="s2">&quot;bottom&quot;</span><span class="p">],</span> <span class="n">skip</span><span class="o">=</span><span class="p">(</span><span class="n">skip_ticks_data_x</span><span class="p">,</span> <span class="n">skip_ticks_data_y</span><span class="p">)</span>
                <span class="p">)</span>
                <span class="n">ax</span><span class="o">.</span><span class="n">set_xlabel</span><span class="p">(</span><span class="sa">r</span><span class="s2">&quot;</span><span class="si">%s</span><span class="s2">&quot;</span> <span class="o">%</span> <span class="n">time_label</span><span class="p">,</span> <span class="n">fontsize</span><span class="o">=</span><span class="n">label_fontsize</span><span class="p">)</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">_make_nice_axes</span><span class="p">(</span><span class="n">ax</span><span class="p">,</span> <span class="n">where</span><span class="o">=</span><span class="p">[</span><span class="s2">&quot;left&quot;</span><span class="p">],</span> <span class="n">skip</span><span class="o">=</span><span class="p">(</span><span class="n">skip_ticks_data_x</span><span class="p">,</span> <span class="n">skip_ticks_data_y</span><span class="p">))</span>
            <span class="c1"># ax.get_xaxis().get_major_formatter().set_useOffset(False)</span>

            <span class="n">ax</span><span class="o">.</span><span class="n">xaxis</span><span class="o">.</span><span class="n">set_major_formatter</span><span class="p">(</span><span class="n">FormatStrFormatter</span><span class="p">(</span><span class="s2">&quot;</span><span class="si">%.0f</span><span class="s2">&quot;</span><span class="p">))</span>
            <span class="n">ax</span><span class="o">.</span><span class="n">label_outer</span><span class="p">()</span>

            <span class="n">ax</span><span class="o">.</span><span class="n">set_xlim</span><span class="p">(</span><span class="n">time</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">time</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">])</span>

            <span class="c1"># trans = transforms.blended_transform_factory(fig.transFigure, ax.transAxes)</span>
            <span class="k">if</span> <span class="n">i</span> <span class="ow">in</span> <span class="n">N_index</span><span class="p">:</span>
                <span class="k">if</span> <span class="n">var_units</span><span class="p">[</span><span class="n">N_index</span><span class="o">.</span><span class="n">index</span><span class="p">(</span><span class="n">i</span><span class="p">)]:</span>
                    <span class="n">ax</span><span class="o">.</span><span class="n">set_ylabel</span><span class="p">(</span><span class="sa">r</span><span class="s2">&quot;</span><span class="si">%s</span><span class="s2"> [</span><span class="si">%s</span><span class="s2">]&quot;</span> <span class="o">%</span> <span class="p">(</span><span class="n">var_names</span><span class="p">[</span><span class="n">N_index</span><span class="o">.</span><span class="n">index</span><span class="p">(</span><span class="n">i</span><span class="p">)],</span> <span class="n">var_units</span><span class="p">[</span><span class="n">N_index</span><span class="o">.</span><span class="n">index</span><span class="p">(</span><span class="n">i</span><span class="p">)]),</span>
                                  <span class="n">fontsize</span><span class="o">=</span><span class="n">label_fontsize</span><span class="p">)</span>
                <span class="k">else</span><span class="p">:</span>
                    <span class="n">ax</span><span class="o">.</span><span class="n">set_ylabel</span><span class="p">(</span><span class="sa">r</span><span class="s2">&quot;</span><span class="si">%s</span><span class="s2">&quot;</span> <span class="o">%</span> <span class="p">(</span><span class="n">var_names</span><span class="p">[</span><span class="n">N_index</span><span class="o">.</span><span class="n">index</span><span class="p">(</span><span class="n">i</span><span class="p">)]),</span> <span class="n">fontsize</span><span class="o">=</span><span class="n">label_fontsize</span><span class="p">)</span>

            <span class="n">ax</span><span class="o">.</span><span class="n">tick_params</span><span class="p">(</span><span class="n">axis</span><span class="o">=</span><span class="s1">&#39;both&#39;</span><span class="p">,</span> <span class="n">which</span><span class="o">=</span><span class="s1">&#39;major&#39;</span><span class="p">,</span> <span class="n">labelsize</span><span class="o">=</span><span class="n">tick_label_size</span><span class="p">)</span>
            <span class="c1"># ax.tick_params(axis=&#39;both&#39;, which=&#39;minor&#39;, labelsize=tick_label_size)</span>

    <span class="n">_add_timeseries</span><span class="p">(</span>
        <span class="n">dataframe</span><span class="o">=</span><span class="n">dataframe</span><span class="p">,</span>
        <span class="n">fig_axes</span><span class="o">=</span><span class="p">(</span><span class="n">fig</span><span class="p">,</span> <span class="n">axes</span><span class="p">),</span>
        <span class="n">grey_masked_samples</span><span class="o">=</span><span class="n">grey_masked_samples</span><span class="p">,</span>
        <span class="n">show_meanline</span><span class="o">=</span><span class="n">show_meanline</span><span class="p">,</span>
        <span class="n">data_linewidth</span><span class="o">=</span><span class="n">data_linewidth</span><span class="p">,</span>
        <span class="n">color</span><span class="o">=</span><span class="n">color</span><span class="p">,</span>
        <span class="n">selected_dataset</span><span class="o">=</span><span class="n">selected_dataset</span><span class="p">,</span>
        <span class="n">alpha</span><span class="o">=</span><span class="n">alpha</span><span class="p">,</span>
        <span class="n">selected_variables</span><span class="o">=</span><span class="n">selected_variables</span>
    <span class="p">)</span>

    <span class="k">if</span> <span class="n">adjust_plot</span><span class="p">:</span>
        <span class="n">fig</span><span class="o">.</span><span class="n">subplots_adjust</span><span class="p">(</span><span class="n">bottom</span><span class="o">=</span><span class="mf">0.15</span><span class="p">,</span> <span class="n">top</span><span class="o">=</span><span class="mf">0.9</span><span class="p">,</span> <span class="n">left</span><span class="o">=</span><span class="mf">0.15</span><span class="p">,</span> <span class="n">right</span><span class="o">=</span><span class="mf">0.95</span><span class="p">,</span> <span class="n">hspace</span><span class="o">=</span><span class="mf">0.3</span><span class="p">)</span>
        <span class="n">pyplot</span><span class="o">.</span><span class="n">tight_layout</span><span class="p">()</span>

    <span class="k">if</span> <span class="n">save_name</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">fig</span><span class="o">.</span><span class="n">savefig</span><span class="p">(</span><span class="n">save_name</span><span class="p">)</span>

    <span class="k">return</span> <span class="n">fig</span><span class="p">,</span> <span class="n">axes</span></div>


<div class="viewcode-block" id="plot_lagfuncs"><a class="viewcode-back" href="../../index.html#tigramite.plotting.plot_lagfuncs">[docs]</a><span class="k">def</span> <span class="nf">plot_lagfuncs</span><span class="p">(</span><span class="n">val_matrix</span><span class="p">,</span> 
                  <span class="n">name</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> 
                  <span class="n">setup_args</span><span class="o">=</span><span class="p">{},</span> 
                  <span class="n">add_lagfunc_args</span><span class="o">=</span><span class="p">{}):</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;Wrapper helper function to plot lag functions.</span>
<span class="sd">    Sets up the matrix object and plots the lagfunction, see parameters in</span>
<span class="sd">    setup_matrix and add_lagfuncs.</span>

<span class="sd">    Parameters</span>
<span class="sd">    ----------</span>
<span class="sd">    val_matrix : array_like</span>
<span class="sd">        Matrix of shape (N, N, tau_max+1) containing test statistic values.</span>
<span class="sd">    name : str, optional (default: None)</span>
<span class="sd">        File name. If None, figure is shown in window.</span>
<span class="sd">    setup_args : dict</span>
<span class="sd">        Arguments for setting up the lag function matrix, see doc of</span>
<span class="sd">        setup_matrix.</span>
<span class="sd">    add_lagfunc_args : dict</span>
<span class="sd">        Arguments for adding a lag function matrix, see doc of add_lagfuncs.</span>

<span class="sd">    Returns</span>
<span class="sd">    -------</span>
<span class="sd">    matrix : object</span>
<span class="sd">        Further lag functions can be overlaid using the</span>
<span class="sd">        matrix.add_lagfuncs(val_matrix) function.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="n">N</span><span class="p">,</span> <span class="n">N</span><span class="p">,</span> <span class="n">tau_max_plusone</span> <span class="o">=</span> <span class="n">val_matrix</span><span class="o">.</span><span class="n">shape</span>
    <span class="n">tau_max</span> <span class="o">=</span> <span class="n">tau_max_plusone</span> <span class="o">-</span> <span class="mi">1</span>

    <span class="n">matrix</span> <span class="o">=</span> <span class="n">setup_matrix</span><span class="p">(</span><span class="n">N</span><span class="o">=</span><span class="n">N</span><span class="p">,</span> <span class="n">tau_max</span><span class="o">=</span><span class="n">tau_max</span><span class="p">,</span> <span class="o">**</span><span class="n">setup_args</span><span class="p">)</span>
    <span class="n">matrix</span><span class="o">.</span><span class="n">add_lagfuncs</span><span class="p">(</span><span class="n">val_matrix</span><span class="o">=</span><span class="n">val_matrix</span><span class="p">,</span> <span class="o">**</span><span class="n">add_lagfunc_args</span><span class="p">)</span>
    <span class="n">matrix</span><span class="o">.</span><span class="n">savefig</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="n">name</span><span class="p">)</span>

    <span class="k">return</span> <span class="n">matrix</span></div>


<div class="viewcode-block" id="setup_matrix"><a class="viewcode-back" href="../../index.html#tigramite.plotting.setup_matrix">[docs]</a><span class="k">class</span> <span class="nc">setup_matrix</span><span class="p">:</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;Create matrix of lag function panels.</span>
<span class="sd">    Class to setup figure object. The function add_lagfuncs(...) allows to plot</span>
<span class="sd">    the val_matrix of shape (N, N, tau_max+1). Multiple lagfunctions can be</span>
<span class="sd">    overlaid for comparison.</span>

<span class="sd">    Parameters</span>
<span class="sd">    ----------</span>
<span class="sd">    N : int</span>
<span class="sd">        Number of variables</span>
<span class="sd">    tau_max : int</span>
<span class="sd">        Maximum time lag.</span>
<span class="sd">    var_names : list, optional (default: None)</span>
<span class="sd">        List of variable names. If None, range(N) is used.</span>
<span class="sd">    figsize : tuple of floats, optional (default: None)</span>
<span class="sd">        Figure size if new figure is created. If None, default pyplot figsize</span>
<span class="sd">        is used.</span>
<span class="sd">    minimum : int, optional (default: -1)</span>
<span class="sd">        Lower y-axis limit.</span>
<span class="sd">    maximum : int, optional (default: 1)</span>
<span class="sd">        Upper y-axis limit.</span>
<span class="sd">    label_space_left : float, optional (default: 0.1)</span>
<span class="sd">        Fraction of horizontal figure space to allocate left of plot for labels.</span>
<span class="sd">    label_space_top : float, optional (default: 0.05)</span>
<span class="sd">        Fraction of vertical figure space to allocate top of plot for labels.</span>
<span class="sd">    legend_width : float, optional (default: 0.15)</span>
<span class="sd">        Fraction of horizontal figure space to allocate right of plot for</span>
<span class="sd">        legend.</span>
<span class="sd">    tick_label_size : int, optional (default: 6)</span>
<span class="sd">        Fontsize of tick labels.</span>
<span class="sd">    x_base : float, optional (default: 1.)</span>
<span class="sd">        x-tick intervals to show.</span>
<span class="sd">    y_base : float, optional (default: .4)</span>
<span class="sd">        y-tick intervals to show.</span>
<span class="sd">    plot_gridlines : bool, optional (default: False)</span>
<span class="sd">        Whether to show a grid.</span>
<span class="sd">    lag_units : str, optional (default: &#39;&#39;)</span>
<span class="sd">    lag_array : array, optional (default: None)</span>
<span class="sd">        Optional specification of lags overwriting np.arange(0, tau_max+1)</span>
<span class="sd">    label_fontsize : int, optional (default: 10)</span>
<span class="sd">        Fontsize of variable labels.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span>
        <span class="bp">self</span><span class="p">,</span>
        <span class="n">N</span><span class="p">,</span>
        <span class="n">tau_max</span><span class="p">,</span>
        <span class="n">var_names</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
        <span class="n">figsize</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
        <span class="n">minimum</span><span class="o">=-</span><span class="mi">1</span><span class="p">,</span>
        <span class="n">maximum</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span>
        <span class="n">label_space_left</span><span class="o">=</span><span class="mf">0.1</span><span class="p">,</span>
        <span class="n">label_space_top</span><span class="o">=</span><span class="mf">0.05</span><span class="p">,</span>
        <span class="n">legend_width</span><span class="o">=</span><span class="mf">0.15</span><span class="p">,</span>
        <span class="n">legend_fontsize</span><span class="o">=</span><span class="mi">10</span><span class="p">,</span>
        <span class="n">x_base</span><span class="o">=</span><span class="mf">1.0</span><span class="p">,</span>
        <span class="n">y_base</span><span class="o">=</span><span class="mf">0.5</span><span class="p">,</span>
        <span class="n">tick_label_size</span><span class="o">=</span><span class="mi">6</span><span class="p">,</span>
        <span class="n">plot_gridlines</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
        <span class="n">lag_units</span><span class="o">=</span><span class="s2">&quot;&quot;</span><span class="p">,</span>
        <span class="n">lag_array</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
        <span class="n">label_fontsize</span><span class="o">=</span><span class="mi">10</span><span class="p">,</span>
    <span class="p">):</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">tau_max</span> <span class="o">=</span> <span class="n">tau_max</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">labels</span> <span class="o">=</span> <span class="p">[]</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">lag_units</span> <span class="o">=</span> <span class="n">lag_units</span>
        <span class="c1"># if lag_array is None:</span>
        <span class="c1">#     self.lag_array = np.arange(0, self.tau_max + 1)</span>
        <span class="c1"># else:</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">lag_array</span> <span class="o">=</span> <span class="n">lag_array</span>
        <span class="k">if</span> <span class="n">x_base</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">x_base</span> <span class="o">=</span> <span class="mi">1</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">x_base</span> <span class="o">=</span> <span class="n">x_base</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">legend_width</span> <span class="o">=</span> <span class="n">legend_width</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">legend_fontsize</span> <span class="o">=</span> <span class="n">legend_fontsize</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">label_space_left</span> <span class="o">=</span> <span class="n">label_space_left</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">label_space_top</span> <span class="o">=</span> <span class="n">label_space_top</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">label_fontsize</span> <span class="o">=</span> <span class="n">label_fontsize</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">fig</span> <span class="o">=</span> <span class="n">pyplot</span><span class="o">.</span><span class="n">figure</span><span class="p">(</span><span class="n">figsize</span><span class="o">=</span><span class="n">figsize</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span> <span class="o">=</span> <span class="p">{}</span>

        <span class="k">if</span> <span class="n">var_names</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">var_names</span> <span class="o">=</span> <span class="nb">range</span><span class="p">(</span><span class="n">N</span><span class="p">)</span>

        <span class="n">plot_index</span> <span class="o">=</span> <span class="mi">1</span>
        <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">N</span><span class="p">):</span>
            <span class="k">for</span> <span class="n">j</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">N</span><span class="p">):</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">[(</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">)]</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">fig</span><span class="o">.</span><span class="n">add_subplot</span><span class="p">(</span><span class="n">N</span><span class="p">,</span> <span class="n">N</span><span class="p">,</span> <span class="n">plot_index</span><span class="p">)</span>
                <span class="c1"># Plot process labels</span>
                <span class="k">if</span> <span class="n">j</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
                    <span class="n">trans</span> <span class="o">=</span> <span class="n">transforms</span><span class="o">.</span><span class="n">blended_transform_factory</span><span class="p">(</span>
                        <span class="bp">self</span><span class="o">.</span><span class="n">fig</span><span class="o">.</span><span class="n">transFigure</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">[(</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">)]</span><span class="o">.</span><span class="n">transAxes</span>
                    <span class="p">)</span>
                    <span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">[(</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">)]</span><span class="o">.</span><span class="n">text</span><span class="p">(</span>
                        <span class="mf">0.01</span><span class="p">,</span>
                        <span class="mf">0.5</span><span class="p">,</span>
                        <span class="s2">&quot;</span><span class="si">%s</span><span class="s2">&quot;</span> <span class="o">%</span> <span class="nb">str</span><span class="p">(</span><span class="n">var_names</span><span class="p">[</span><span class="n">i</span><span class="p">]),</span>
                        <span class="n">fontsize</span><span class="o">=</span><span class="n">label_fontsize</span><span class="p">,</span>
                        <span class="n">horizontalalignment</span><span class="o">=</span><span class="s2">&quot;left&quot;</span><span class="p">,</span>
                        <span class="n">verticalalignment</span><span class="o">=</span><span class="s2">&quot;center&quot;</span><span class="p">,</span>
                        <span class="n">transform</span><span class="o">=</span><span class="n">trans</span><span class="p">,</span>
                    <span class="p">)</span>
                <span class="k">if</span> <span class="n">i</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
                    <span class="n">trans</span> <span class="o">=</span> <span class="n">transforms</span><span class="o">.</span><span class="n">blended_transform_factory</span><span class="p">(</span>
                        <span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">[(</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">)]</span><span class="o">.</span><span class="n">transAxes</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">fig</span><span class="o">.</span><span class="n">transFigure</span>
                    <span class="p">)</span>
                    <span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">[(</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">)]</span><span class="o">.</span><span class="n">text</span><span class="p">(</span>
                        <span class="mf">0.5</span><span class="p">,</span>
                        <span class="mf">0.99</span><span class="p">,</span>
                        <span class="sa">r</span><span class="s2">&quot;${\to}$ &quot;</span> <span class="o">+</span> <span class="s2">&quot;</span><span class="si">%s</span><span class="s2">&quot;</span> <span class="o">%</span> <span class="nb">str</span><span class="p">(</span><span class="n">var_names</span><span class="p">[</span><span class="n">j</span><span class="p">]),</span>
                        <span class="n">fontsize</span><span class="o">=</span><span class="n">label_fontsize</span><span class="p">,</span>
                        <span class="n">horizontalalignment</span><span class="o">=</span><span class="s2">&quot;center&quot;</span><span class="p">,</span>
                        <span class="n">verticalalignment</span><span class="o">=</span><span class="s2">&quot;top&quot;</span><span class="p">,</span>
                        <span class="n">transform</span><span class="o">=</span><span class="n">trans</span><span class="p">,</span>
                    <span class="p">)</span>

                <span class="c1"># Make nice axis</span>
                <span class="n">_make_nice_axes</span><span class="p">(</span>
                    <span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">[(</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">)],</span> <span class="n">where</span><span class="o">=</span><span class="p">[</span><span class="s2">&quot;left&quot;</span><span class="p">,</span> <span class="s2">&quot;bottom&quot;</span><span class="p">],</span> <span class="n">skip</span><span class="o">=</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span>
                <span class="p">)</span>
                <span class="k">if</span> <span class="n">x_base</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                    <span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">[(</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">)]</span><span class="o">.</span><span class="n">xaxis</span><span class="o">.</span><span class="n">set_major_locator</span><span class="p">(</span>
                        <span class="n">ticker</span><span class="o">.</span><span class="n">FixedLocator</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">tau_max</span> <span class="o">+</span> <span class="mi">1</span><span class="p">,</span> <span class="n">x_base</span><span class="p">))</span>
                    <span class="p">)</span>
                    <span class="k">if</span> <span class="n">x_base</span> <span class="o">/</span> <span class="mf">2.0</span> <span class="o">%</span> <span class="mi">1</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
                        <span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">[(</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">)]</span><span class="o">.</span><span class="n">xaxis</span><span class="o">.</span><span class="n">set_minor_locator</span><span class="p">(</span>
                            <span class="n">ticker</span><span class="o">.</span><span class="n">FixedLocator</span><span class="p">(</span>
                                <span class="n">np</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">tau_max</span> <span class="o">+</span> <span class="mi">1</span><span class="p">,</span> <span class="n">x_base</span> <span class="o">/</span> <span class="mf">2.0</span><span class="p">)</span>
                            <span class="p">)</span>
                        <span class="p">)</span>
                <span class="k">if</span> <span class="n">y_base</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                    <span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">[(</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">)]</span><span class="o">.</span><span class="n">yaxis</span><span class="o">.</span><span class="n">set_major_locator</span><span class="p">(</span>
                        <span class="n">ticker</span><span class="o">.</span><span class="n">FixedLocator</span><span class="p">(</span>
                            <span class="n">np</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span>
                                <span class="n">_myround</span><span class="p">(</span><span class="n">minimum</span><span class="p">,</span> <span class="n">y_base</span><span class="p">,</span> <span class="s2">&quot;down&quot;</span><span class="p">),</span>
                                <span class="n">_myround</span><span class="p">(</span><span class="n">maximum</span><span class="p">,</span> <span class="n">y_base</span><span class="p">,</span> <span class="s2">&quot;up&quot;</span><span class="p">)</span> <span class="o">+</span> <span class="n">y_base</span><span class="p">,</span>
                                <span class="n">y_base</span><span class="p">,</span>
                            <span class="p">)</span>
                        <span class="p">)</span>
                    <span class="p">)</span>
                    <span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">[(</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">)]</span><span class="o">.</span><span class="n">yaxis</span><span class="o">.</span><span class="n">set_minor_locator</span><span class="p">(</span>
                        <span class="n">ticker</span><span class="o">.</span><span class="n">FixedLocator</span><span class="p">(</span>
                            <span class="n">np</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span>
                                <span class="n">_myround</span><span class="p">(</span><span class="n">minimum</span><span class="p">,</span> <span class="n">y_base</span><span class="p">,</span> <span class="s2">&quot;down&quot;</span><span class="p">),</span>
                                <span class="n">_myround</span><span class="p">(</span><span class="n">maximum</span><span class="p">,</span> <span class="n">y_base</span><span class="p">,</span> <span class="s2">&quot;up&quot;</span><span class="p">)</span> <span class="o">+</span> <span class="n">y_base</span><span class="p">,</span>
                                <span class="n">y_base</span> <span class="o">/</span> <span class="mf">2.0</span><span class="p">,</span>
                            <span class="p">)</span>
                        <span class="p">)</span>
                    <span class="p">)</span>

                    <span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">[(</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">)]</span><span class="o">.</span><span class="n">set_ylim</span><span class="p">(</span>
                        <span class="n">_myround</span><span class="p">(</span><span class="n">minimum</span><span class="p">,</span> <span class="n">y_base</span><span class="p">,</span> <span class="s2">&quot;down&quot;</span><span class="p">),</span>
                        <span class="n">_myround</span><span class="p">(</span><span class="n">maximum</span><span class="p">,</span> <span class="n">y_base</span><span class="p">,</span> <span class="s2">&quot;up&quot;</span><span class="p">),</span>
                    <span class="p">)</span>
                <span class="k">if</span> <span class="n">j</span> <span class="o">!=</span> <span class="mi">0</span><span class="p">:</span>
                    <span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">[(</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">)]</span><span class="o">.</span><span class="n">get_yaxis</span><span class="p">()</span><span class="o">.</span><span class="n">set_ticklabels</span><span class="p">([])</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">[(</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">)]</span><span class="o">.</span><span class="n">set_xlim</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">tau_max</span><span class="p">)</span>
                <span class="k">if</span> <span class="n">plot_gridlines</span><span class="p">:</span>
                    <span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">[(</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">)]</span><span class="o">.</span><span class="n">grid</span><span class="p">(</span>
                        <span class="kc">True</span><span class="p">,</span>
                        <span class="n">which</span><span class="o">=</span><span class="s2">&quot;major&quot;</span><span class="p">,</span>
                        <span class="n">color</span><span class="o">=</span><span class="s2">&quot;black&quot;</span><span class="p">,</span>
                        <span class="n">linestyle</span><span class="o">=</span><span class="s2">&quot;dotted&quot;</span><span class="p">,</span>
                        <span class="n">dashes</span><span class="o">=</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">),</span>
                        <span class="n">linewidth</span><span class="o">=</span><span class="mf">0.05</span><span class="p">,</span>
                        <span class="n">zorder</span><span class="o">=-</span><span class="mi">5</span><span class="p">,</span>
                    <span class="p">)</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">[(</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">)]</span><span class="o">.</span><span class="n">tick_params</span><span class="p">(</span><span class="n">axis</span><span class="o">=</span><span class="s1">&#39;both&#39;</span><span class="p">,</span> <span class="n">which</span><span class="o">=</span><span class="s1">&#39;major&#39;</span><span class="p">,</span> <span class="n">labelsize</span><span class="o">=</span><span class="n">tick_label_size</span><span class="p">)</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">[(</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">)]</span><span class="o">.</span><span class="n">tick_params</span><span class="p">(</span><span class="n">axis</span><span class="o">=</span><span class="s1">&#39;both&#39;</span><span class="p">,</span> <span class="n">which</span><span class="o">=</span><span class="s1">&#39;minor&#39;</span><span class="p">,</span> <span class="n">labelsize</span><span class="o">=</span><span class="n">tick_label_size</span><span class="p">)</span>

                <span class="n">plot_index</span> <span class="o">+=</span> <span class="mi">1</span>

<div class="viewcode-block" id="setup_matrix.add_lagfuncs"><a class="viewcode-back" href="../../index.html#tigramite.plotting.setup_matrix.add_lagfuncs">[docs]</a>    <span class="k">def</span> <span class="nf">add_lagfuncs</span><span class="p">(</span>
        <span class="bp">self</span><span class="p">,</span>
        <span class="n">val_matrix</span><span class="p">,</span>
        <span class="n">sig_thres</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
        <span class="n">conf_matrix</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
        <span class="n">color</span><span class="o">=</span><span class="s2">&quot;black&quot;</span><span class="p">,</span>
        <span class="n">label</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
        <span class="n">two_sided_thres</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
        <span class="n">marker</span><span class="o">=</span><span class="s2">&quot;.&quot;</span><span class="p">,</span>
        <span class="n">markersize</span><span class="o">=</span><span class="mi">5</span><span class="p">,</span>
        <span class="n">alpha</span><span class="o">=</span><span class="mf">1.0</span><span class="p">,</span>
    <span class="p">):</span>
<span class="w">        </span><span class="sd">&quot;&quot;&quot;Add lag function plot from val_matrix array.</span>

<span class="sd">        Parameters</span>
<span class="sd">        ----------</span>
<span class="sd">        val_matrix : array_like</span>
<span class="sd">            Matrix of shape (N, N, tau_max+1) containing test statistic values.</span>
<span class="sd">        sig_thres : array-like, optional (default: None)</span>
<span class="sd">            Matrix of significance thresholds. Must be of same shape as</span>
<span class="sd">            val_matrix.</span>
<span class="sd">        conf_matrix : array-like, optional (default: None)</span>
<span class="sd">            Matrix of shape (, N, tau_max+1, 2) containing confidence bounds.</span>
<span class="sd">        color : str, optional (default: &#39;black&#39;)</span>
<span class="sd">            Line color.</span>
<span class="sd">        label : str</span>
<span class="sd">            Test statistic label.</span>
<span class="sd">        two_sided_thres : bool, optional (default: True)</span>
<span class="sd">            Whether to draw sig_thres for pos. and neg. values.</span>
<span class="sd">        marker : matplotlib marker symbol, optional (default: &#39;.&#39;)</span>
<span class="sd">            Marker.</span>
<span class="sd">        markersize : int, optional (default: 5)</span>
<span class="sd">            Marker size.</span>
<span class="sd">        alpha : float, optional (default: 1.)</span>
<span class="sd">            Opacity.</span>
<span class="sd">        &quot;&quot;&quot;</span>

        <span class="k">if</span> <span class="n">label</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">labels</span><span class="o">.</span><span class="n">append</span><span class="p">((</span><span class="n">label</span><span class="p">,</span> <span class="n">color</span><span class="p">,</span> <span class="n">marker</span><span class="p">,</span> <span class="n">markersize</span><span class="p">,</span> <span class="n">alpha</span><span class="p">))</span>

        <span class="k">for</span> <span class="n">ij</span> <span class="ow">in</span> <span class="nb">list</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">):</span>
            <span class="n">i</span> <span class="o">=</span> <span class="n">ij</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
            <span class="n">j</span> <span class="o">=</span> <span class="n">ij</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span>
            <span class="n">maskedres</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">copy</span><span class="p">(</span><span class="n">val_matrix</span><span class="p">[</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">,</span> <span class="nb">int</span><span class="p">(</span><span class="n">i</span> <span class="o">==</span> <span class="n">j</span><span class="p">)</span> <span class="p">:])</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">[(</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">)]</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span>
                <span class="nb">range</span><span class="p">(</span><span class="nb">int</span><span class="p">(</span><span class="n">i</span> <span class="o">==</span> <span class="n">j</span><span class="p">),</span> <span class="bp">self</span><span class="o">.</span><span class="n">tau_max</span> <span class="o">+</span> <span class="mi">1</span><span class="p">),</span>
                <span class="n">maskedres</span><span class="p">,</span>
                <span class="n">linestyle</span><span class="o">=</span><span class="s2">&quot;&quot;</span><span class="p">,</span>
                <span class="n">color</span><span class="o">=</span><span class="n">color</span><span class="p">,</span>
                <span class="n">marker</span><span class="o">=</span><span class="n">marker</span><span class="p">,</span>
                <span class="n">markersize</span><span class="o">=</span><span class="n">markersize</span><span class="p">,</span>
                <span class="n">alpha</span><span class="o">=</span><span class="n">alpha</span><span class="p">,</span>
                <span class="n">clip_on</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
            <span class="p">)</span>
            <span class="k">if</span> <span class="n">conf_matrix</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                <span class="n">maskedconfres</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">copy</span><span class="p">(</span><span class="n">conf_matrix</span><span class="p">[</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">,</span> <span class="nb">int</span><span class="p">(</span><span class="n">i</span> <span class="o">==</span> <span class="n">j</span><span class="p">)</span> <span class="p">:])</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">[(</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">)]</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span>
                    <span class="nb">range</span><span class="p">(</span><span class="nb">int</span><span class="p">(</span><span class="n">i</span> <span class="o">==</span> <span class="n">j</span><span class="p">),</span> <span class="bp">self</span><span class="o">.</span><span class="n">tau_max</span> <span class="o">+</span> <span class="mi">1</span><span class="p">),</span>
                    <span class="n">maskedconfres</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">],</span>
                    <span class="n">linestyle</span><span class="o">=</span><span class="s2">&quot;&quot;</span><span class="p">,</span>
                    <span class="n">color</span><span class="o">=</span><span class="n">color</span><span class="p">,</span>
                    <span class="n">marker</span><span class="o">=</span><span class="s2">&quot;_&quot;</span><span class="p">,</span>
                    <span class="n">markersize</span><span class="o">=</span><span class="n">markersize</span> <span class="o">-</span> <span class="mi">2</span><span class="p">,</span>
                    <span class="n">alpha</span><span class="o">=</span><span class="n">alpha</span><span class="p">,</span>
                    <span class="n">clip_on</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
                <span class="p">)</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">[(</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">)]</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span>
                    <span class="nb">range</span><span class="p">(</span><span class="nb">int</span><span class="p">(</span><span class="n">i</span> <span class="o">==</span> <span class="n">j</span><span class="p">),</span> <span class="bp">self</span><span class="o">.</span><span class="n">tau_max</span> <span class="o">+</span> <span class="mi">1</span><span class="p">),</span>
                    <span class="n">maskedconfres</span><span class="p">[:,</span> <span class="mi">1</span><span class="p">],</span>
                    <span class="n">linestyle</span><span class="o">=</span><span class="s2">&quot;&quot;</span><span class="p">,</span>
                    <span class="n">color</span><span class="o">=</span><span class="n">color</span><span class="p">,</span>
                    <span class="n">marker</span><span class="o">=</span><span class="s2">&quot;_&quot;</span><span class="p">,</span>
                    <span class="n">markersize</span><span class="o">=</span><span class="n">markersize</span> <span class="o">-</span> <span class="mi">2</span><span class="p">,</span>
                    <span class="n">alpha</span><span class="o">=</span><span class="n">alpha</span><span class="p">,</span>
                    <span class="n">clip_on</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
                <span class="p">)</span>

            <span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">[(</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">)]</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span>
                <span class="nb">range</span><span class="p">(</span><span class="nb">int</span><span class="p">(</span><span class="n">i</span> <span class="o">==</span> <span class="n">j</span><span class="p">),</span> <span class="bp">self</span><span class="o">.</span><span class="n">tau_max</span> <span class="o">+</span> <span class="mi">1</span><span class="p">),</span>
                <span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">tau_max</span> <span class="o">+</span> <span class="mi">1</span> <span class="o">-</span> <span class="nb">int</span><span class="p">(</span><span class="n">i</span> <span class="o">==</span> <span class="n">j</span><span class="p">)),</span>
                <span class="n">color</span><span class="o">=</span><span class="s2">&quot;black&quot;</span><span class="p">,</span>
                <span class="n">linestyle</span><span class="o">=</span><span class="s2">&quot;dotted&quot;</span><span class="p">,</span>
                <span class="n">linewidth</span><span class="o">=</span><span class="mf">0.1</span><span class="p">,</span>
            <span class="p">)</span>

            <span class="k">if</span> <span class="n">sig_thres</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                <span class="n">maskedsigres</span> <span class="o">=</span> <span class="n">sig_thres</span><span class="p">[</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">,</span> <span class="nb">int</span><span class="p">(</span><span class="n">i</span> <span class="o">==</span> <span class="n">j</span><span class="p">)</span> <span class="p">:]</span>

                <span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">[(</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">)]</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span>
                    <span class="nb">range</span><span class="p">(</span><span class="nb">int</span><span class="p">(</span><span class="n">i</span> <span class="o">==</span> <span class="n">j</span><span class="p">),</span> <span class="bp">self</span><span class="o">.</span><span class="n">tau_max</span> <span class="o">+</span> <span class="mi">1</span><span class="p">),</span>
                    <span class="n">maskedsigres</span><span class="p">,</span>
                    <span class="n">color</span><span class="o">=</span><span class="n">color</span><span class="p">,</span>
                    <span class="n">linestyle</span><span class="o">=</span><span class="s2">&quot;solid&quot;</span><span class="p">,</span>
                    <span class="n">linewidth</span><span class="o">=</span><span class="mf">0.1</span><span class="p">,</span>
                    <span class="n">alpha</span><span class="o">=</span><span class="n">alpha</span><span class="p">,</span>
                <span class="p">)</span>
                <span class="k">if</span> <span class="n">two_sided_thres</span><span class="p">:</span>
                    <span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">[(</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">)]</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span>
                        <span class="nb">range</span><span class="p">(</span><span class="nb">int</span><span class="p">(</span><span class="n">i</span> <span class="o">==</span> <span class="n">j</span><span class="p">),</span> <span class="bp">self</span><span class="o">.</span><span class="n">tau_max</span> <span class="o">+</span> <span class="mi">1</span><span class="p">),</span>
                        <span class="o">-</span><span class="n">sig_thres</span><span class="p">[</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">,</span> <span class="nb">int</span><span class="p">(</span><span class="n">i</span> <span class="o">==</span> <span class="n">j</span><span class="p">)</span> <span class="p">:],</span>
                        <span class="n">color</span><span class="o">=</span><span class="n">color</span><span class="p">,</span>
                        <span class="n">linestyle</span><span class="o">=</span><span class="s2">&quot;solid&quot;</span><span class="p">,</span>
                        <span class="n">linewidth</span><span class="o">=</span><span class="mf">0.1</span><span class="p">,</span>
                        <span class="n">alpha</span><span class="o">=</span><span class="n">alpha</span><span class="p">,</span>
                    <span class="p">)</span></div>
        <span class="c1"># pyplot.tight_layout()</span>

<div class="viewcode-block" id="setup_matrix.savefig"><a class="viewcode-back" href="../../index.html#tigramite.plotting.setup_matrix.savefig">[docs]</a>    <span class="k">def</span> <span class="nf">savefig</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">name</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
<span class="w">        </span><span class="sd">&quot;&quot;&quot;Save matrix figure.</span>

<span class="sd">        Parameters</span>
<span class="sd">        ----------</span>
<span class="sd">        name : str, optional (default: None)</span>
<span class="sd">            File name. If None, figure is shown in window.</span>
<span class="sd">        &quot;&quot;&quot;</span>

        <span class="c1"># Trick to plot legend</span>
        <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">labels</span><span class="p">)</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
            <span class="n">axlegend</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">fig</span><span class="o">.</span><span class="n">add_subplot</span><span class="p">(</span><span class="mi">111</span><span class="p">,</span> <span class="n">frameon</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span>
            <span class="n">axlegend</span><span class="o">.</span><span class="n">spines</span><span class="p">[</span><span class="s2">&quot;left&quot;</span><span class="p">]</span><span class="o">.</span><span class="n">set_color</span><span class="p">(</span><span class="s2">&quot;none&quot;</span><span class="p">)</span>
            <span class="n">axlegend</span><span class="o">.</span><span class="n">spines</span><span class="p">[</span><span class="s2">&quot;right&quot;</span><span class="p">]</span><span class="o">.</span><span class="n">set_color</span><span class="p">(</span><span class="s2">&quot;none&quot;</span><span class="p">)</span>
            <span class="n">axlegend</span><span class="o">.</span><span class="n">spines</span><span class="p">[</span><span class="s2">&quot;bottom&quot;</span><span class="p">]</span><span class="o">.</span><span class="n">set_color</span><span class="p">(</span><span class="s2">&quot;none&quot;</span><span class="p">)</span>
            <span class="n">axlegend</span><span class="o">.</span><span class="n">spines</span><span class="p">[</span><span class="s2">&quot;top&quot;</span><span class="p">]</span><span class="o">.</span><span class="n">set_color</span><span class="p">(</span><span class="s2">&quot;none&quot;</span><span class="p">)</span>
            <span class="n">axlegend</span><span class="o">.</span><span class="n">set_xticks</span><span class="p">([])</span>
            <span class="n">axlegend</span><span class="o">.</span><span class="n">set_yticks</span><span class="p">([])</span>

            <span class="c1"># self.labels.append((label, color, marker, markersize, alpha))</span>
            <span class="k">for</span> <span class="n">item</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">labels</span><span class="p">:</span>
                <span class="n">label</span> <span class="o">=</span> <span class="n">item</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
                <span class="n">color</span> <span class="o">=</span> <span class="n">item</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span>
                <span class="n">marker</span> <span class="o">=</span> <span class="n">item</span><span class="p">[</span><span class="mi">2</span><span class="p">]</span>
                <span class="n">markersize</span> <span class="o">=</span> <span class="n">item</span><span class="p">[</span><span class="mi">3</span><span class="p">]</span>
                <span class="n">alpha</span> <span class="o">=</span> <span class="n">item</span><span class="p">[</span><span class="mi">4</span><span class="p">]</span>

                <span class="n">axlegend</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span>
                    <span class="p">[],</span>
                    <span class="p">[],</span>
                    <span class="n">linestyle</span><span class="o">=</span><span class="s2">&quot;&quot;</span><span class="p">,</span>
                    <span class="n">color</span><span class="o">=</span><span class="n">color</span><span class="p">,</span>
                    <span class="n">marker</span><span class="o">=</span><span class="n">marker</span><span class="p">,</span>
                    <span class="n">markersize</span><span class="o">=</span><span class="n">markersize</span><span class="p">,</span>
                    <span class="n">label</span><span class="o">=</span><span class="n">label</span><span class="p">,</span>
                    <span class="n">alpha</span><span class="o">=</span><span class="n">alpha</span><span class="p">,</span>
                <span class="p">)</span>
            <span class="n">axlegend</span><span class="o">.</span><span class="n">legend</span><span class="p">(</span>
                <span class="n">loc</span><span class="o">=</span><span class="s2">&quot;upper left&quot;</span><span class="p">,</span>
                <span class="n">ncol</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span>
                <span class="n">bbox_to_anchor</span><span class="o">=</span><span class="p">(</span><span class="mf">1.05</span><span class="p">,</span> <span class="mf">0.0</span><span class="p">,</span> <span class="mf">0.1</span><span class="p">,</span> <span class="mf">1.0</span><span class="p">),</span>
                <span class="n">borderaxespad</span><span class="o">=</span><span class="mi">0</span><span class="p">,</span>
                <span class="n">fontsize</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">legend_fontsize</span><span class="p">,</span>
            <span class="p">)</span><span class="o">.</span><span class="n">draw_frame</span><span class="p">(</span><span class="kc">False</span><span class="p">)</span>

            <span class="bp">self</span><span class="o">.</span><span class="n">fig</span><span class="o">.</span><span class="n">subplots_adjust</span><span class="p">(</span>
                <span class="n">left</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">label_space_left</span><span class="p">,</span>
                <span class="n">right</span><span class="o">=</span><span class="mf">1.0</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">legend_width</span><span class="p">,</span>
                <span class="n">top</span><span class="o">=</span><span class="mf">1.0</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">label_space_top</span><span class="p">,</span>
                <span class="n">hspace</span><span class="o">=</span><span class="mf">0.35</span><span class="p">,</span>
                <span class="n">wspace</span><span class="o">=</span><span class="mf">0.35</span><span class="p">,</span>
            <span class="p">)</span>
            <span class="n">pyplot</span><span class="o">.</span><span class="n">figtext</span><span class="p">(</span>
                <span class="mf">0.5</span><span class="p">,</span>
                <span class="mf">0.01</span><span class="p">,</span>
                <span class="sa">r</span><span class="s2">&quot;lag $\tau$ [</span><span class="si">%s</span><span class="s2">]&quot;</span> <span class="o">%</span> <span class="bp">self</span><span class="o">.</span><span class="n">lag_units</span><span class="p">,</span>
                <span class="n">horizontalalignment</span><span class="o">=</span><span class="s2">&quot;center&quot;</span><span class="p">,</span>
                <span class="n">fontsize</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">label_fontsize</span><span class="p">,</span>
            <span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">fig</span><span class="o">.</span><span class="n">subplots_adjust</span><span class="p">(</span>
                <span class="n">left</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">label_space_left</span><span class="p">,</span>
                <span class="n">right</span><span class="o">=</span><span class="mf">0.95</span><span class="p">,</span>
                <span class="n">top</span><span class="o">=</span><span class="mf">1.0</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">label_space_top</span><span class="p">,</span>
                <span class="n">hspace</span><span class="o">=</span><span class="mf">0.35</span><span class="p">,</span>
                <span class="n">wspace</span><span class="o">=</span><span class="mf">0.35</span><span class="p">,</span>
            <span class="p">)</span>
            <span class="n">pyplot</span><span class="o">.</span><span class="n">figtext</span><span class="p">(</span>
                <span class="mf">0.55</span><span class="p">,</span>
                <span class="mf">0.01</span><span class="p">,</span>
                <span class="sa">r</span><span class="s2">&quot;lag $\tau$ [</span><span class="si">%s</span><span class="s2">]&quot;</span> <span class="o">%</span> <span class="bp">self</span><span class="o">.</span><span class="n">lag_units</span><span class="p">,</span>
                <span class="n">horizontalalignment</span><span class="o">=</span><span class="s2">&quot;center&quot;</span><span class="p">,</span>
                <span class="n">fontsize</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">label_fontsize</span><span class="p">,</span>
            <span class="p">)</span>

        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">lag_array</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="k">assert</span> <span class="bp">self</span><span class="o">.</span><span class="n">lag_array</span><span class="o">.</span><span class="n">shape</span> <span class="o">==</span> <span class="n">np</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">tau_max</span> <span class="o">+</span> <span class="mi">1</span><span class="p">)</span><span class="o">.</span><span class="n">shape</span>
            <span class="k">for</span> <span class="n">ij</span> <span class="ow">in</span> <span class="nb">list</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">):</span>
                <span class="n">i</span> <span class="o">=</span> <span class="n">ij</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
                <span class="n">j</span> <span class="o">=</span> <span class="n">ij</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">[(</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">)]</span><span class="o">.</span><span class="n">set_xticklabels</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">lag_array</span><span class="p">[::</span> <span class="bp">self</span><span class="o">.</span><span class="n">x_base</span><span class="p">])</span>

        <span class="k">if</span> <span class="n">name</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">fig</span><span class="o">.</span><span class="n">savefig</span><span class="p">(</span><span class="n">name</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">pyplot</span><span class="o">.</span><span class="n">show</span><span class="p">()</span></div></div>



<div class="viewcode-block" id="plot_scatterplots"><a class="viewcode-back" href="../../index.html#tigramite.plotting.plot_scatterplots">[docs]</a><span class="k">def</span> <span class="nf">plot_scatterplots</span><span class="p">(</span><span class="n">dataframe</span><span class="p">,</span> 
                      <span class="n">name</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> 
                      <span class="n">setup_args</span><span class="o">=</span><span class="p">{},</span> 
                      <span class="n">add_scatterplot_args</span><span class="o">=</span><span class="p">{},</span>
                      <span class="n">selected_dataset</span><span class="o">=</span><span class="mi">0</span><span class="p">):</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;Wrapper helper function to plot scatter plots.</span>
<span class="sd">    Sets up the matrix object and plots the scatter plots, see parameters in</span>
<span class="sd">    setup_scatter_matrix and add_scatterplot.</span>

<span class="sd">    Parameters</span>
<span class="sd">    ----------</span>
<span class="sd">    dataframe : data object</span>
<span class="sd">        Tigramite dataframe object. It must have the attributes dataframe.values</span>
<span class="sd">        yielding a numpy array of shape (observations T, variables N) and</span>
<span class="sd">        optionally a mask of the same shape and a missing values flag.</span>
<span class="sd">    name : str, optional (default: None)</span>
<span class="sd">        File name. If None, figure is shown in window.</span>
<span class="sd">    setup_args : dict</span>
<span class="sd">        Arguments for setting up the scatter plot matrix, see doc of</span>
<span class="sd">        setup_scatter_matrix.</span>
<span class="sd">    add_scatterplot_args : dict</span>
<span class="sd">        Arguments for adding a scatter plot matrix.</span>
<span class="sd">    selected_dataset : int, optional (default: 0)</span>
<span class="sd">        In case of multiple datasets in dataframe, plot this one.</span>

<span class="sd">    Returns</span>
<span class="sd">    -------</span>
<span class="sd">    matrix : object</span>
<span class="sd">        Further scatter plot can be overlaid using the</span>
<span class="sd">        matrix.add_scatterplot function.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="n">N</span> <span class="o">=</span> <span class="n">dataframe</span><span class="o">.</span><span class="n">N</span>

    <span class="n">matrix</span> <span class="o">=</span> <span class="n">setup_scatter_matrix</span><span class="p">(</span><span class="n">N</span><span class="o">=</span><span class="n">N</span><span class="p">,</span> <span class="n">var_names</span><span class="o">=</span><span class="n">dataframe</span><span class="o">.</span><span class="n">var_names</span><span class="p">,</span> <span class="o">**</span><span class="n">setup_args</span><span class="p">)</span>
    <span class="n">matrix</span><span class="o">.</span><span class="n">add_scatterplot</span><span class="p">(</span><span class="n">dataframe</span><span class="o">=</span><span class="n">dataframe</span><span class="p">,</span> <span class="n">selected_dataset</span><span class="o">=</span><span class="n">selected_dataset</span><span class="p">,</span> <span class="o">**</span><span class="n">add_scatterplot_args</span><span class="p">)</span>
    <span class="n">matrix</span><span class="o">.</span><span class="n">adjustfig</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="n">name</span><span class="p">)</span>
   

    <span class="k">return</span> <span class="n">matrix</span></div>


<div class="viewcode-block" id="setup_scatter_matrix"><a class="viewcode-back" href="../../index.html#tigramite.plotting.setup_scatter_matrix">[docs]</a><span class="k">class</span> <span class="nc">setup_scatter_matrix</span><span class="p">:</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;Create matrix of scatter plot panels.</span>
<span class="sd">    Class to setup figure object. The function add_scatterplot allows to plot</span>
<span class="sd">    scatterplots of variables in the dataframe. Multiple scatter plots can be</span>
<span class="sd">    overlaid for comparison.</span>

<span class="sd">    Parameters</span>
<span class="sd">    ----------</span>
<span class="sd">    N : int</span>
<span class="sd">        Number of variables</span>
<span class="sd">    var_names : list, optional (default: None)</span>
<span class="sd">        List of variable names. If None, range(N) is used.</span>
<span class="sd">    figsize : tuple of floats, optional (default: None)</span>
<span class="sd">        Figure size if new figure is created. If None, default pyplot figsize</span>
<span class="sd">        is used.</span>
<span class="sd">    label_space_left : float, optional (default: 0.1)</span>
<span class="sd">        Fraction of horizontal figure space to allocate left of plot for labels.</span>
<span class="sd">    label_space_top : float, optional (default: 0.05)</span>
<span class="sd">        Fraction of vertical figure space to allocate top of plot for labels.</span>
<span class="sd">    legend_width : float, optional (default: 0.15)</span>
<span class="sd">        Fraction of horizontal figure space to allocate right of plot for</span>
<span class="sd">        legend.</span>
<span class="sd">    tick_label_size : int, optional (default: 6)</span>
<span class="sd">        Fontsize of tick labels.</span>
<span class="sd">    plot_gridlines : bool, optional (default: False)</span>
<span class="sd">        Whether to show a grid.</span>
<span class="sd">    label_fontsize : int, optional (default: 10)</span>
<span class="sd">        Fontsize of variable labels.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span>
        <span class="bp">self</span><span class="p">,</span>
        <span class="n">N</span><span class="p">,</span>
        <span class="n">var_names</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
        <span class="n">figsize</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
        <span class="n">label_space_left</span><span class="o">=</span><span class="mf">0.1</span><span class="p">,</span>
        <span class="n">label_space_top</span><span class="o">=</span><span class="mf">0.05</span><span class="p">,</span>
        <span class="n">legend_width</span><span class="o">=</span><span class="mf">0.15</span><span class="p">,</span>
        <span class="n">legend_fontsize</span><span class="o">=</span><span class="mi">10</span><span class="p">,</span>
        <span class="n">plot_gridlines</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
        <span class="n">tick_label_size</span><span class="o">=</span><span class="mi">6</span><span class="p">,</span>
        <span class="n">label_fontsize</span><span class="o">=</span><span class="mi">10</span><span class="p">,</span>
    <span class="p">):</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">labels</span> <span class="o">=</span> <span class="p">[]</span>
    
        <span class="bp">self</span><span class="o">.</span><span class="n">legend_width</span> <span class="o">=</span> <span class="n">legend_width</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">legend_fontsize</span> <span class="o">=</span> <span class="n">legend_fontsize</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">label_space_left</span> <span class="o">=</span> <span class="n">label_space_left</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">label_space_top</span> <span class="o">=</span> <span class="n">label_space_top</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">label_fontsize</span> <span class="o">=</span> <span class="n">label_fontsize</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">fig</span> <span class="o">=</span> <span class="n">pyplot</span><span class="o">.</span><span class="n">figure</span><span class="p">(</span><span class="n">figsize</span><span class="o">=</span><span class="n">figsize</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span> <span class="o">=</span> <span class="p">{}</span>

        <span class="k">if</span> <span class="n">var_names</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">var_names</span> <span class="o">=</span> <span class="nb">range</span><span class="p">(</span><span class="n">N</span><span class="p">)</span>

        <span class="n">plot_index</span> <span class="o">=</span> <span class="mi">1</span>
        <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">N</span><span class="p">):</span>
            <span class="k">for</span> <span class="n">j</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">N</span><span class="p">):</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">[(</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">)]</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">fig</span><span class="o">.</span><span class="n">add_subplot</span><span class="p">(</span><span class="n">N</span><span class="p">,</span> <span class="n">N</span><span class="p">,</span> <span class="n">plot_index</span><span class="p">,</span> <span class="n">axes_class</span><span class="o">=</span><span class="n">Axes</span><span class="p">)</span>
                <span class="c1"># Plot process labels</span>
                <span class="k">if</span> <span class="n">j</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
                    <span class="n">trans</span> <span class="o">=</span> <span class="n">transforms</span><span class="o">.</span><span class="n">blended_transform_factory</span><span class="p">(</span>
                        <span class="bp">self</span><span class="o">.</span><span class="n">fig</span><span class="o">.</span><span class="n">transFigure</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">[(</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">)]</span><span class="o">.</span><span class="n">transAxes</span>
                    <span class="p">)</span>
                    <span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">[(</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">)]</span><span class="o">.</span><span class="n">text</span><span class="p">(</span>
                        <span class="mf">0.01</span><span class="p">,</span>
                        <span class="mf">0.5</span><span class="p">,</span>
                        <span class="s2">&quot;</span><span class="si">%s</span><span class="s2">&quot;</span> <span class="o">%</span> <span class="nb">str</span><span class="p">(</span><span class="n">var_names</span><span class="p">[</span><span class="n">i</span><span class="p">]),</span>
                        <span class="n">fontsize</span><span class="o">=</span><span class="n">label_fontsize</span><span class="p">,</span>
                        <span class="n">horizontalalignment</span><span class="o">=</span><span class="s2">&quot;left&quot;</span><span class="p">,</span>
                        <span class="n">verticalalignment</span><span class="o">=</span><span class="s2">&quot;center&quot;</span><span class="p">,</span>
                        <span class="n">transform</span><span class="o">=</span><span class="n">trans</span><span class="p">,</span>
                    <span class="p">)</span>
                <span class="k">if</span> <span class="n">i</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
                    <span class="n">trans</span> <span class="o">=</span> <span class="n">transforms</span><span class="o">.</span><span class="n">blended_transform_factory</span><span class="p">(</span>
                        <span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">[(</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">)]</span><span class="o">.</span><span class="n">transAxes</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">fig</span><span class="o">.</span><span class="n">transFigure</span>
                    <span class="p">)</span>
                    <span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">[(</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">)]</span><span class="o">.</span><span class="n">text</span><span class="p">(</span>
                        <span class="mf">0.5</span><span class="p">,</span>
                        <span class="mf">0.99</span><span class="p">,</span>
                        <span class="sa">r</span><span class="s2">&quot;${\to}$ &quot;</span> <span class="o">+</span> <span class="s2">&quot;</span><span class="si">%s</span><span class="s2">&quot;</span> <span class="o">%</span> <span class="nb">str</span><span class="p">(</span><span class="n">var_names</span><span class="p">[</span><span class="n">j</span><span class="p">]),</span>
                        <span class="n">fontsize</span><span class="o">=</span><span class="n">label_fontsize</span><span class="p">,</span>
                        <span class="n">horizontalalignment</span><span class="o">=</span><span class="s2">&quot;center&quot;</span><span class="p">,</span>
                        <span class="n">verticalalignment</span><span class="o">=</span><span class="s2">&quot;top&quot;</span><span class="p">,</span>
                        <span class="n">transform</span><span class="o">=</span><span class="n">trans</span><span class="p">,</span>
                    <span class="p">)</span>

                <span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">[(</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">)]</span><span class="o">.</span><span class="n">axis</span><span class="p">[</span><span class="s2">&quot;right&quot;</span><span class="p">]</span><span class="o">.</span><span class="n">set_visible</span><span class="p">(</span><span class="kc">False</span><span class="p">)</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">[(</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">)]</span><span class="o">.</span><span class="n">axis</span><span class="p">[</span><span class="s2">&quot;top&quot;</span><span class="p">]</span><span class="o">.</span><span class="n">set_visible</span><span class="p">(</span><span class="kc">False</span><span class="p">)</span>

                <span class="k">if</span> <span class="n">j</span> <span class="o">!=</span> <span class="mi">0</span><span class="p">:</span>
                    <span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">[(</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">)]</span><span class="o">.</span><span class="n">get_yaxis</span><span class="p">()</span><span class="o">.</span><span class="n">set_ticklabels</span><span class="p">([])</span>
                <span class="k">if</span> <span class="n">i</span> <span class="o">!=</span> <span class="n">N</span> <span class="o">-</span> <span class="mi">1</span><span class="p">:</span>
                    <span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">[(</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">)]</span><span class="o">.</span><span class="n">get_xaxis</span><span class="p">()</span><span class="o">.</span><span class="n">set_ticklabels</span><span class="p">([])</span>

                <span class="k">if</span> <span class="n">plot_gridlines</span><span class="p">:</span>
                    <span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">[(</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">)]</span><span class="o">.</span><span class="n">grid</span><span class="p">(</span>
                        <span class="kc">True</span><span class="p">,</span>
                        <span class="n">which</span><span class="o">=</span><span class="s2">&quot;major&quot;</span><span class="p">,</span>
                        <span class="n">color</span><span class="o">=</span><span class="s2">&quot;black&quot;</span><span class="p">,</span>
                        <span class="n">linestyle</span><span class="o">=</span><span class="s2">&quot;dotted&quot;</span><span class="p">,</span>
                        <span class="n">dashes</span><span class="o">=</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">),</span>
                        <span class="n">linewidth</span><span class="o">=</span><span class="mf">0.05</span><span class="p">,</span>
                        <span class="n">zorder</span><span class="o">=-</span><span class="mi">5</span><span class="p">,</span>
                    <span class="p">)</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">[(</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">)]</span><span class="o">.</span><span class="n">tick_params</span><span class="p">(</span><span class="n">axis</span><span class="o">=</span><span class="s1">&#39;both&#39;</span><span class="p">,</span> <span class="n">which</span><span class="o">=</span><span class="s1">&#39;major&#39;</span><span class="p">,</span> <span class="n">labelsize</span><span class="o">=</span><span class="n">tick_label_size</span><span class="p">)</span>

                <span class="n">plot_index</span> <span class="o">+=</span> <span class="mi">1</span>

<div class="viewcode-block" id="setup_scatter_matrix.add_scatterplot"><a class="viewcode-back" href="../../index.html#tigramite.plotting.setup_scatter_matrix.add_scatterplot">[docs]</a>    <span class="k">def</span> <span class="nf">add_scatterplot</span><span class="p">(</span>
        <span class="bp">self</span><span class="p">,</span>
        <span class="n">dataframe</span><span class="p">,</span>
        <span class="n">matrix_lags</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
        <span class="n">color</span><span class="o">=</span><span class="s2">&quot;black&quot;</span><span class="p">,</span>
        <span class="n">label</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
        <span class="n">marker</span><span class="o">=</span><span class="s2">&quot;.&quot;</span><span class="p">,</span>
        <span class="n">markersize</span><span class="o">=</span><span class="mi">5</span><span class="p">,</span>
        <span class="n">alpha</span><span class="o">=</span><span class="mf">.2</span><span class="p">,</span>
        <span class="n">selected_dataset</span><span class="o">=</span><span class="mi">0</span><span class="p">,</span>
    <span class="p">):</span>
<span class="w">        </span><span class="sd">&quot;&quot;&quot;Add scatter plot.</span>

<span class="sd">        Parameters</span>
<span class="sd">        ----------</span>
<span class="sd">        dataframe : data object</span>
<span class="sd">            Tigramite dataframe object. It must have the attributes dataframe.values</span>
<span class="sd">            yielding a numpy array of shape (observations T, variables N) and</span>
<span class="sd">            optionally a mask of the same shape and a missing values flag.</span>
<span class="sd">        matrix_lags : array</span>
<span class="sd">            Lags to use in scatter plots. Either None or of shape (N, N). Then the</span>
<span class="sd">            entry matrix_lags[i, j] = tau will depict the scatter plot of </span>
<span class="sd">            time series (i, -tau) vs (j, 0). If None, tau = 0 for i != j and for i = j</span>
<span class="sd">            tau = 1. </span>
<span class="sd">        color : str, optional (default: &#39;black&#39;)</span>
<span class="sd">            Line color.</span>
<span class="sd">        label : str</span>
<span class="sd">            Test statistic label.</span>
<span class="sd">        marker : matplotlib marker symbol, optional (default: &#39;.&#39;)</span>
<span class="sd">            Marker.</span>
<span class="sd">        markersize : int, optional (default: 5)</span>
<span class="sd">            Marker size.</span>
<span class="sd">        alpha : float, optional (default: 1.)</span>
<span class="sd">            Opacity.</span>
<span class="sd">        selected_dataset : int, optional (default: 0)</span>
<span class="sd">            In case of multiple datasets in dataframe, plot this one.</span>
<span class="sd">        &quot;&quot;&quot;</span>

        <span class="k">if</span> <span class="n">matrix_lags</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> <span class="n">np</span><span class="o">.</span><span class="n">any</span><span class="p">(</span><span class="n">matrix_lags</span> <span class="o">&lt;</span> <span class="mi">0</span><span class="p">):</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;matrix_lags must be non-negative!&quot;</span><span class="p">)</span>

        <span class="n">data</span> <span class="o">=</span> <span class="n">dataframe</span><span class="o">.</span><span class="n">values</span><span class="p">[</span><span class="n">selected_dataset</span><span class="p">]</span>
        <span class="k">if</span> <span class="n">dataframe</span><span class="o">.</span><span class="n">mask</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">mask</span> <span class="o">=</span> <span class="n">dataframe</span><span class="o">.</span><span class="n">mask</span><span class="p">[</span><span class="n">selected_dataset</span><span class="p">]</span>

        <span class="n">T</span><span class="p">,</span> <span class="n">dim</span> <span class="o">=</span> <span class="n">data</span><span class="o">.</span><span class="n">shape</span>

        <span class="k">if</span> <span class="n">label</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">labels</span><span class="o">.</span><span class="n">append</span><span class="p">((</span><span class="n">label</span><span class="p">,</span> <span class="n">color</span><span class="p">,</span> <span class="n">marker</span><span class="p">,</span> <span class="n">markersize</span><span class="p">,</span> <span class="n">alpha</span><span class="p">))</span>

        <span class="k">for</span> <span class="n">ij</span> <span class="ow">in</span> <span class="nb">list</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">):</span>                
            <span class="n">i</span> <span class="o">=</span> <span class="n">ij</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
            <span class="n">j</span> <span class="o">=</span> <span class="n">ij</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span>
            <span class="k">if</span> <span class="n">matrix_lags</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
                <span class="k">if</span> <span class="n">i</span> <span class="o">==</span> <span class="n">j</span><span class="p">:</span>
                    <span class="n">lag</span> <span class="o">=</span> <span class="mi">1</span>
                <span class="k">else</span><span class="p">:</span>
                    <span class="n">lag</span> <span class="o">=</span> <span class="mi">0</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">lag</span> <span class="o">=</span> <span class="n">matrix_lags</span><span class="p">[</span><span class="n">i</span><span class="p">,</span><span class="n">j</span><span class="p">]</span>
            <span class="n">x</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">copy</span><span class="p">(</span><span class="n">data</span><span class="p">[:</span><span class="n">T</span><span class="o">-</span><span class="n">lag</span><span class="p">,</span> <span class="n">i</span><span class="p">])</span>
            <span class="n">y</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">copy</span><span class="p">(</span><span class="n">data</span><span class="p">[</span><span class="n">lag</span><span class="p">:,</span> <span class="n">j</span><span class="p">])</span>
            <span class="k">if</span> <span class="n">dataframe</span><span class="o">.</span><span class="n">mask</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                <span class="n">x</span><span class="p">[</span><span class="n">mask</span><span class="p">[:</span><span class="n">T</span><span class="o">-</span><span class="n">lag</span><span class="p">,</span> <span class="n">i</span><span class="p">]</span><span class="o">==</span><span class="mi">1</span><span class="p">]</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">nan</span>
                <span class="n">y</span><span class="p">[</span><span class="n">mask</span><span class="p">[</span><span class="n">lag</span><span class="p">:,</span> <span class="n">j</span><span class="p">]</span><span class="o">==</span><span class="mi">1</span><span class="p">]</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">nan</span>

            <span class="c1"># print(i, j, lag, x.shape, y.shape)</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">[(</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">)]</span><span class="o">.</span><span class="n">scatter</span><span class="p">(</span>
                <span class="n">x</span><span class="p">,</span> <span class="n">y</span><span class="p">,</span>
                <span class="n">color</span><span class="o">=</span><span class="n">color</span><span class="p">,</span>
                <span class="n">marker</span><span class="o">=</span><span class="n">marker</span><span class="p">,</span>
                <span class="n">s</span><span class="o">=</span><span class="n">markersize</span><span class="p">,</span>
                <span class="n">alpha</span><span class="o">=</span><span class="n">alpha</span><span class="p">,</span>
                <span class="n">clip_on</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
                <span class="n">label</span><span class="o">=</span><span class="sa">r</span><span class="s2">&quot;$\tau{=}</span><span class="si">%d</span><span class="s2">$&quot;</span> <span class="o">%</span><span class="n">lag</span><span class="p">,</span>
            <span class="p">)</span></div>
            <span class="c1"># self.axes_dict[(i, j)].text(0., 1., r&quot;$\tau{=}%d$&quot; %lag, </span>
            <span class="c1">#     fontsize=self.legend_fontsize,</span>
            <span class="c1">#     ha=&#39;left&#39;, va=&#39;top&#39;,</span>
            <span class="c1">#     transform=self.axes_dict[(i, j)].transAxes)</span>


<div class="viewcode-block" id="setup_scatter_matrix.adjustfig"><a class="viewcode-back" href="../../index.html#tigramite.plotting.setup_scatter_matrix.adjustfig">[docs]</a>    <span class="k">def</span> <span class="nf">adjustfig</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">name</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
<span class="w">        </span><span class="sd">&quot;&quot;&quot;Adjust matrix figure.</span>

<span class="sd">        Parameters</span>
<span class="sd">        ----------</span>
<span class="sd">        name : str, optional (default: None)</span>
<span class="sd">            File name. If None, figure is shown in window.</span>
<span class="sd">        &quot;&quot;&quot;</span>

        <span class="c1"># Trick to plot legends</span>
        <span class="n">colors</span> <span class="o">=</span> <span class="p">[]</span>
        <span class="k">for</span> <span class="n">item</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">labels</span><span class="p">:</span>
            <span class="n">colors</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">item</span><span class="p">[</span><span class="mi">1</span><span class="p">])</span>
        <span class="k">for</span> <span class="n">ij</span> <span class="ow">in</span> <span class="nb">list</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">):</span>                
            <span class="n">i</span> <span class="o">=</span> <span class="n">ij</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
            <span class="n">j</span> <span class="o">=</span> <span class="n">ij</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span>

            <span class="n">leg</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">[(</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">)]</span><span class="o">.</span><span class="n">legend</span><span class="p">(</span>
                <span class="c1"># loc=&quot;upper left&quot;,</span>
                <span class="n">ncol</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span>
                <span class="c1"># bbox_to_anchor=(1.05, 0.0, 0.1, 1.0),</span>
                <span class="c1"># borderaxespad=0,</span>
                <span class="n">fontsize</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">legend_fontsize</span><span class="o">-</span><span class="mi">2</span><span class="p">,</span>
                <span class="n">labelcolor</span><span class="o">=</span><span class="n">colors</span><span class="p">,</span>
                <span class="p">)</span><span class="o">.</span><span class="n">draw_frame</span><span class="p">(</span><span class="kc">False</span><span class="p">)</span>
        
        <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">labels</span><span class="p">)</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
            <span class="n">axlegend</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">fig</span><span class="o">.</span><span class="n">add_subplot</span><span class="p">(</span><span class="mi">111</span><span class="p">,</span> <span class="n">frameon</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span>
            <span class="n">axlegend</span><span class="o">.</span><span class="n">spines</span><span class="p">[</span><span class="s2">&quot;left&quot;</span><span class="p">]</span><span class="o">.</span><span class="n">set_color</span><span class="p">(</span><span class="s2">&quot;none&quot;</span><span class="p">)</span>
            <span class="n">axlegend</span><span class="o">.</span><span class="n">spines</span><span class="p">[</span><span class="s2">&quot;right&quot;</span><span class="p">]</span><span class="o">.</span><span class="n">set_color</span><span class="p">(</span><span class="s2">&quot;none&quot;</span><span class="p">)</span>
            <span class="n">axlegend</span><span class="o">.</span><span class="n">spines</span><span class="p">[</span><span class="s2">&quot;bottom&quot;</span><span class="p">]</span><span class="o">.</span><span class="n">set_color</span><span class="p">(</span><span class="s2">&quot;none&quot;</span><span class="p">)</span>
            <span class="n">axlegend</span><span class="o">.</span><span class="n">spines</span><span class="p">[</span><span class="s2">&quot;top&quot;</span><span class="p">]</span><span class="o">.</span><span class="n">set_color</span><span class="p">(</span><span class="s2">&quot;none&quot;</span><span class="p">)</span>
            <span class="n">axlegend</span><span class="o">.</span><span class="n">set_xticks</span><span class="p">([])</span>
            <span class="n">axlegend</span><span class="o">.</span><span class="n">set_yticks</span><span class="p">([])</span>

            <span class="c1"># self.labels.append((label, color, marker, markersize, alpha))</span>
            <span class="k">for</span> <span class="n">item</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">labels</span><span class="p">:</span>
                <span class="n">label</span> <span class="o">=</span> <span class="n">item</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
                <span class="n">color</span> <span class="o">=</span> <span class="n">item</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span>
                <span class="n">marker</span> <span class="o">=</span> <span class="n">item</span><span class="p">[</span><span class="mi">2</span><span class="p">]</span>
                <span class="n">markersize</span> <span class="o">=</span> <span class="n">item</span><span class="p">[</span><span class="mi">3</span><span class="p">]</span>
                <span class="n">alpha</span> <span class="o">=</span> <span class="n">item</span><span class="p">[</span><span class="mi">4</span><span class="p">]</span>

                <span class="n">axlegend</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span>
                    <span class="p">[],</span>
                    <span class="p">[],</span>
                    <span class="n">linestyle</span><span class="o">=</span><span class="s2">&quot;&quot;</span><span class="p">,</span>
                    <span class="n">color</span><span class="o">=</span><span class="n">color</span><span class="p">,</span>
                    <span class="n">marker</span><span class="o">=</span><span class="n">marker</span><span class="p">,</span>
                    <span class="n">markersize</span><span class="o">=</span><span class="n">markersize</span><span class="p">,</span>
                    <span class="n">label</span><span class="o">=</span><span class="n">label</span><span class="p">,</span>
                    <span class="n">alpha</span><span class="o">=</span><span class="n">alpha</span><span class="p">,</span>
                <span class="p">)</span>
            <span class="n">axlegend</span><span class="o">.</span><span class="n">legend</span><span class="p">(</span>
                <span class="n">loc</span><span class="o">=</span><span class="s2">&quot;upper left&quot;</span><span class="p">,</span>
                <span class="n">ncol</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span>
                <span class="n">bbox_to_anchor</span><span class="o">=</span><span class="p">(</span><span class="mf">1.05</span><span class="p">,</span> <span class="mf">0.0</span><span class="p">,</span> <span class="mf">0.1</span><span class="p">,</span> <span class="mf">1.0</span><span class="p">),</span>
                <span class="n">borderaxespad</span><span class="o">=</span><span class="mi">0</span><span class="p">,</span>
                <span class="n">fontsize</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">legend_fontsize</span><span class="p">,</span>
            <span class="p">)</span><span class="o">.</span><span class="n">draw_frame</span><span class="p">(</span><span class="kc">False</span><span class="p">)</span>

            <span class="bp">self</span><span class="o">.</span><span class="n">fig</span><span class="o">.</span><span class="n">subplots_adjust</span><span class="p">(</span>
                <span class="n">bottom</span><span class="o">=</span><span class="mf">0.05</span><span class="p">,</span>
                <span class="n">left</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">label_space_left</span><span class="p">,</span>
                <span class="n">right</span><span class="o">=</span><span class="mf">1.0</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">legend_width</span><span class="p">,</span>
                <span class="n">top</span><span class="o">=</span><span class="mf">1.0</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">label_space_top</span><span class="p">,</span>
                <span class="n">hspace</span><span class="o">=</span><span class="mf">0.5</span><span class="p">,</span>
                <span class="n">wspace</span><span class="o">=</span><span class="mf">0.35</span><span class="p">,</span>
            <span class="p">)</span>
      
        <span class="k">else</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">fig</span><span class="o">.</span><span class="n">subplots_adjust</span><span class="p">(</span>
                <span class="n">left</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">label_space_left</span><span class="p">,</span>
                <span class="n">bottom</span><span class="o">=</span><span class="mf">0.05</span><span class="p">,</span>
                <span class="n">right</span><span class="o">=</span><span class="mf">0.95</span><span class="p">,</span>
                <span class="n">top</span><span class="o">=</span><span class="mf">1.0</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">label_space_top</span><span class="p">,</span>
                <span class="n">hspace</span><span class="o">=</span><span class="mf">0.35</span><span class="p">,</span>
                <span class="n">wspace</span><span class="o">=</span><span class="mf">0.35</span><span class="p">,</span>
            <span class="p">)</span>
       
        <span class="k">if</span> <span class="n">name</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">fig</span><span class="o">.</span><span class="n">savefig</span><span class="p">(</span><span class="n">name</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">pyplot</span><span class="o">.</span><span class="n">show</span><span class="p">()</span></div></div>


<div class="viewcode-block" id="plot_densityplots"><a class="viewcode-back" href="../../index.html#tigramite.plotting.plot_densityplots">[docs]</a><span class="k">def</span> <span class="nf">plot_densityplots</span><span class="p">(</span><span class="n">dataframe</span><span class="p">,</span> 
                      <span class="n">name</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> 
                      <span class="n">setup_args</span><span class="o">=</span><span class="p">{},</span> 
                      <span class="n">add_densityplot_args</span><span class="o">=</span><span class="p">{},</span>
                      <span class="n">selected_dataset</span><span class="o">=</span><span class="mi">0</span><span class="p">,</span>
                      <span class="n">show_marginal_densities_on_diagonal</span><span class="o">=</span><span class="kc">True</span><span class="p">):</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;Wrapper helper function to plot density plots.</span>
<span class="sd">    Sets up the matrix object and plots the density plots, see parameters in</span>
<span class="sd">    setup_density_matrix and add_densityplot.</span>

<span class="sd">    The diagonal shows the marginal densities. </span>

<span class="sd">    Requires seaborn.</span>

<span class="sd">    Parameters</span>
<span class="sd">    ----------</span>
<span class="sd">    dataframe : data object</span>
<span class="sd">        Tigramite dataframe object. It must have the attributes dataframe.values</span>
<span class="sd">        yielding a numpy array of shape (observations T, variables N) and</span>
<span class="sd">        optionally a mask of the same shape and a missing values flag.</span>
<span class="sd">    name : str, optional (default: None)</span>
<span class="sd">        File name. If None, figure is shown in window.</span>
<span class="sd">    setup_args : dict</span>
<span class="sd">        Arguments for setting up the density plot matrix, see doc of</span>
<span class="sd">        setup_density_matrix.</span>
<span class="sd">    add_densityplot_args : dict</span>
<span class="sd">        Arguments for adding a density plot matrix.</span>
<span class="sd">    selected_dataset : int, optional (default: 0)</span>
<span class="sd">        In case of multiple datasets in dataframe, plot this one.</span>
<span class="sd">    show_marginal_densities_on_diagonal : bool, optional (default: True)</span>
<span class="sd">        Flag to show marginal densities on the diagonal of the density plots</span>

<span class="sd">    Returns</span>
<span class="sd">    -------</span>
<span class="sd">    matrix : object</span>
<span class="sd">        Further density plots can be overlaid using the</span>
<span class="sd">        matrix.add_densityplot function.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="n">N</span> <span class="o">=</span> <span class="n">dataframe</span><span class="o">.</span><span class="n">N</span>

    <span class="n">matrix</span> <span class="o">=</span> <span class="n">setup_density_matrix</span><span class="p">(</span><span class="n">N</span><span class="o">=</span><span class="n">N</span><span class="p">,</span> <span class="n">var_names</span><span class="o">=</span><span class="n">dataframe</span><span class="o">.</span><span class="n">var_names</span><span class="p">,</span> <span class="o">**</span><span class="n">setup_args</span><span class="p">)</span>
    <span class="n">matrix</span><span class="o">.</span><span class="n">add_densityplot</span><span class="p">(</span><span class="n">dataframe</span><span class="o">=</span><span class="n">dataframe</span><span class="p">,</span> <span class="n">selected_dataset</span><span class="o">=</span><span class="n">selected_dataset</span><span class="p">,</span>
                           <span class="n">show_marginal_densities_on_diagonal</span><span class="o">=</span><span class="n">show_marginal_densities_on_diagonal</span><span class="p">,</span> <span class="o">**</span><span class="n">add_densityplot_args</span><span class="p">)</span>
    <span class="n">matrix</span><span class="o">.</span><span class="n">adjustfig</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="n">name</span><span class="p">)</span>
   

    <span class="k">return</span> <span class="n">matrix</span></div>


<div class="viewcode-block" id="setup_density_matrix"><a class="viewcode-back" href="../../index.html#tigramite.plotting.setup_density_matrix">[docs]</a><span class="k">class</span> <span class="nc">setup_density_matrix</span><span class="p">:</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;Create matrix of density plot panels.</span>
<span class="sd">    Class to setup figure object. The function add_densityplot allows to plot</span>
<span class="sd">    density plots of variables in the dataframe. </span>

<span class="sd">    Further density plots can be overlaid using the matrix.add_densityplot </span>
<span class="sd">    function.</span>

<span class="sd">    Parameters</span>
<span class="sd">    ----------</span>
<span class="sd">    N : int</span>
<span class="sd">        Number of variables</span>
<span class="sd">    var_names : list, optional (default: None)</span>
<span class="sd">        List of variable names. If None, range(N) is used.</span>
<span class="sd">    figsize : tuple of floats, optional (default: None)</span>
<span class="sd">        Figure size if new figure is created. If None, default pyplot figsize</span>
<span class="sd">        is used.</span>
<span class="sd">    label_space_left : float, optional (default: 0.1)</span>
<span class="sd">        Fraction of horizontal figure space to allocate left of plot for labels.</span>
<span class="sd">    label_space_top : float, optional (default: 0.05)</span>
<span class="sd">        Fraction of vertical figure space to allocate top of plot for labels.</span>
<span class="sd">    legend_width : float, optional (default: 0.15)</span>
<span class="sd">        Fraction of horizontal figure space to allocate right of plot for</span>
<span class="sd">        legend.</span>
<span class="sd">    tick_label_size : int, optional (default: 6)</span>
<span class="sd">        Fontsize of tick labels.</span>
<span class="sd">    plot_gridlines : bool, optional (default: False)</span>
<span class="sd">        Whether to show a grid.</span>
<span class="sd">    label_fontsize : int, optional (default: 10)</span>
<span class="sd">        Fontsize of variable labels.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span>
        <span class="bp">self</span><span class="p">,</span>
        <span class="n">N</span><span class="p">,</span>
        <span class="n">var_names</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
        <span class="n">figsize</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
        <span class="n">label_space_left</span><span class="o">=</span><span class="mf">0.15</span><span class="p">,</span>
        <span class="n">label_space_top</span><span class="o">=</span><span class="mf">0.05</span><span class="p">,</span>
        <span class="n">legend_width</span><span class="o">=</span><span class="mf">0.15</span><span class="p">,</span>
        <span class="n">legend_fontsize</span><span class="o">=</span><span class="mi">10</span><span class="p">,</span>
        <span class="n">tick_label_size</span><span class="o">=</span><span class="mi">6</span><span class="p">,</span>
        <span class="n">plot_gridlines</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
        <span class="n">label_fontsize</span><span class="o">=</span><span class="mi">10</span><span class="p">,</span>
    <span class="p">):</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">labels</span> <span class="o">=</span> <span class="p">[]</span>
    
        <span class="bp">self</span><span class="o">.</span><span class="n">legend_width</span> <span class="o">=</span> <span class="n">legend_width</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">legend_fontsize</span> <span class="o">=</span> <span class="n">legend_fontsize</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">label_space_left</span> <span class="o">=</span> <span class="n">label_space_left</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">label_space_top</span> <span class="o">=</span> <span class="n">label_space_top</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">label_fontsize</span> <span class="o">=</span> <span class="n">label_fontsize</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">fig</span> <span class="o">=</span> <span class="n">pyplot</span><span class="o">.</span><span class="n">figure</span><span class="p">(</span><span class="n">figsize</span><span class="o">=</span><span class="n">figsize</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span> <span class="o">=</span> <span class="p">{}</span>

        <span class="k">if</span> <span class="n">var_names</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">var_names</span> <span class="o">=</span> <span class="nb">range</span><span class="p">(</span><span class="n">N</span><span class="p">)</span>

        <span class="n">plot_index</span> <span class="o">=</span> <span class="mi">1</span>
        <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">N</span><span class="p">):</span>
            <span class="k">for</span> <span class="n">j</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">N</span><span class="p">):</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">[(</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">)]</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">fig</span><span class="o">.</span><span class="n">add_subplot</span><span class="p">(</span><span class="n">N</span><span class="p">,</span> <span class="n">N</span><span class="p">,</span> <span class="n">plot_index</span><span class="p">)</span>
                <span class="c1"># Plot process labels</span>
                <span class="k">if</span> <span class="n">j</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
                    <span class="n">trans</span> <span class="o">=</span> <span class="n">transforms</span><span class="o">.</span><span class="n">blended_transform_factory</span><span class="p">(</span>
                        <span class="bp">self</span><span class="o">.</span><span class="n">fig</span><span class="o">.</span><span class="n">transFigure</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">[(</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">)]</span><span class="o">.</span><span class="n">transAxes</span>
                    <span class="p">)</span>
                    <span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">[(</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">)]</span><span class="o">.</span><span class="n">text</span><span class="p">(</span>
                        <span class="mf">0.01</span><span class="p">,</span>
                        <span class="mf">0.5</span><span class="p">,</span>
                        <span class="s2">&quot;</span><span class="si">%s</span><span class="s2">&quot;</span> <span class="o">%</span> <span class="nb">str</span><span class="p">(</span><span class="n">var_names</span><span class="p">[</span><span class="n">i</span><span class="p">]),</span>
                        <span class="n">fontsize</span><span class="o">=</span><span class="n">label_fontsize</span><span class="p">,</span>
                        <span class="n">horizontalalignment</span><span class="o">=</span><span class="s2">&quot;left&quot;</span><span class="p">,</span>
                        <span class="n">verticalalignment</span><span class="o">=</span><span class="s2">&quot;center&quot;</span><span class="p">,</span>
                        <span class="n">transform</span><span class="o">=</span><span class="n">trans</span><span class="p">,</span>
                    <span class="p">)</span>
                <span class="k">if</span> <span class="n">i</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
                    <span class="n">trans</span> <span class="o">=</span> <span class="n">transforms</span><span class="o">.</span><span class="n">blended_transform_factory</span><span class="p">(</span>
                        <span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">[(</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">)]</span><span class="o">.</span><span class="n">transAxes</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">fig</span><span class="o">.</span><span class="n">transFigure</span>
                    <span class="p">)</span>
                    <span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">[(</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">)]</span><span class="o">.</span><span class="n">text</span><span class="p">(</span>
                        <span class="mf">0.5</span><span class="p">,</span>
                        <span class="mf">0.99</span><span class="p">,</span>
                        <span class="sa">r</span><span class="s2">&quot;${\to}$ &quot;</span> <span class="o">+</span> <span class="s2">&quot;</span><span class="si">%s</span><span class="s2">&quot;</span> <span class="o">%</span> <span class="nb">str</span><span class="p">(</span><span class="n">var_names</span><span class="p">[</span><span class="n">j</span><span class="p">]),</span>
                        <span class="n">fontsize</span><span class="o">=</span><span class="n">label_fontsize</span><span class="p">,</span>
                        <span class="n">horizontalalignment</span><span class="o">=</span><span class="s2">&quot;center&quot;</span><span class="p">,</span>
                        <span class="n">verticalalignment</span><span class="o">=</span><span class="s2">&quot;top&quot;</span><span class="p">,</span>
                        <span class="n">transform</span><span class="o">=</span><span class="n">trans</span><span class="p">,</span>
                    <span class="p">)</span>

                <span class="c1"># _make_nice_axes(self.axes_dict[(i, j)], where=[&quot;bottom&quot;], skip=(1, 1)             )</span>
                <span class="c1"># self.axes_dict[(i, j)].axis[&quot;right&quot;].set_visible(False)</span>
                <span class="c1"># self.axes_dict[(i, j)].axis[&quot;top&quot;].set_visible(False)</span>
                <span class="k">if</span> <span class="n">i</span> <span class="o">==</span> <span class="n">j</span><span class="p">:</span>
                    <span class="c1"># self.axes_dict[(i, j)].axis[&quot;left&quot;].set_visible(False)</span>
                    <span class="n">_make_nice_axes</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">[(</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">)],</span> <span class="n">where</span><span class="o">=</span><span class="p">[</span><span class="s2">&quot;bottom&quot;</span><span class="p">],</span> <span class="n">skip</span><span class="o">=</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">))</span>
                <span class="k">else</span><span class="p">:</span>
                    <span class="n">_make_nice_axes</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">[(</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">)],</span> <span class="n">where</span><span class="o">=</span><span class="p">[</span><span class="s2">&quot;left&quot;</span><span class="p">,</span> <span class="s2">&quot;bottom&quot;</span><span class="p">],</span> <span class="n">skip</span><span class="o">=</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">))</span>
                <span class="c1"># if j != 0:</span>
                <span class="c1">#     self.axes_dict[(i, j)].get_yaxis().set_ticklabels([])</span>
                <span class="c1"># if i != N - 1:</span>
                <span class="c1">#     self.axes_dict[(i, j)].get_xaxis().set_ticklabels([])</span>

                <span class="k">if</span> <span class="n">plot_gridlines</span><span class="p">:</span>
                    <span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">[(</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">)]</span><span class="o">.</span><span class="n">grid</span><span class="p">(</span>
                        <span class="kc">True</span><span class="p">,</span>
                        <span class="n">which</span><span class="o">=</span><span class="s2">&quot;major&quot;</span><span class="p">,</span>
                        <span class="n">color</span><span class="o">=</span><span class="s2">&quot;black&quot;</span><span class="p">,</span>
                        <span class="n">linestyle</span><span class="o">=</span><span class="s2">&quot;dotted&quot;</span><span class="p">,</span>
                        <span class="n">dashes</span><span class="o">=</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">),</span>
                        <span class="n">linewidth</span><span class="o">=</span><span class="mf">0.05</span><span class="p">,</span>
                        <span class="n">zorder</span><span class="o">=-</span><span class="mi">5</span><span class="p">,</span>
                    <span class="p">)</span>
                <span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">[(</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">)]</span><span class="o">.</span><span class="n">tick_params</span><span class="p">(</span><span class="n">axis</span><span class="o">=</span><span class="s1">&#39;both&#39;</span><span class="p">,</span> <span class="n">which</span><span class="o">=</span><span class="s1">&#39;major&#39;</span><span class="p">,</span> <span class="n">labelsize</span><span class="o">=</span><span class="n">tick_label_size</span><span class="p">)</span>
                <span class="n">plot_index</span> <span class="o">+=</span> <span class="mi">1</span>

<div class="viewcode-block" id="setup_density_matrix.add_densityplot"><a class="viewcode-back" href="../../index.html#tigramite.plotting.setup_density_matrix.add_densityplot">[docs]</a>    <span class="k">def</span> <span class="nf">add_densityplot</span><span class="p">(</span>
        <span class="bp">self</span><span class="p">,</span>
        <span class="n">dataframe</span><span class="p">,</span>
        <span class="n">matrix_lags</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
        <span class="n">label</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
        <span class="n">label_color</span><span class="o">=</span><span class="s1">&#39;black&#39;</span><span class="p">,</span>
        <span class="n">snskdeplot_args</span> <span class="o">=</span> <span class="p">{</span><span class="s1">&#39;cmap&#39;</span><span class="p">:</span><span class="s1">&#39;Greys&#39;</span><span class="p">},</span>
        <span class="n">snskdeplot_diagonal_args</span> <span class="o">=</span> <span class="p">{},</span>
        <span class="n">selected_dataset</span><span class="o">=</span><span class="mi">0</span><span class="p">,</span>
        <span class="n">show_marginal_densities_on_diagonal</span><span class="o">=</span><span class="kc">True</span>
    <span class="p">):</span>
<span class="w">        </span><span class="sd">&quot;&quot;&quot;Add density function plot.</span>

<span class="sd">        Parameters</span>
<span class="sd">        ----------</span>
<span class="sd">        dataframe : data object</span>
<span class="sd">            Tigramite dataframe object. It must have the attributes dataframe.values</span>
<span class="sd">            yielding a numpy array of shape (observations T, variables N) and</span>
<span class="sd">            optionally a mask of the same shape and a missing values flag.</span>
<span class="sd">        matrix_lags : array</span>
<span class="sd">            Lags to use in scatter plots. Either None or non-neg array of shape (N, N). Then the</span>
<span class="sd">            entry matrix_lags[i, j] = tau will depict the scatter plot of </span>
<span class="sd">            time series (i, -tau) vs (j, 0). If None, tau = 0 for i != j and for i = j</span>
<span class="sd">            tau = 1. </span>
<span class="sd">        snskdeplot_args : dict</span>
<span class="sd">            Optional parameters to pass to sns.kdeplot() for i != j for off-diagonal plots.</span>
<span class="sd">        snskdeplot_diagonal_args : dict</span>
<span class="sd">            Optional parameters to pass to sns.kdeplot() for i == j on diagonal.</span>
<span class="sd">        label : string</span>
<span class="sd">            Label of this plot.</span>
<span class="sd">        label_color : string</span>
<span class="sd">            Color of line created just for legend.</span>
<span class="sd">        selected_dataset : int, optional (default: 0)</span>
<span class="sd">            In case of multiple datasets in dataframe, plot this one.</span>
<span class="sd">        show_marginal_densities_on_diagonal : bool, optional (default: True)</span>
<span class="sd">            Flag to show marginal densities on the diagonal of the density plots</span>
<span class="sd">        &quot;&quot;&quot;</span>

        <span class="c1"># Use seaborn for this one</span>
        <span class="kn">import</span> <span class="nn">seaborn</span> <span class="k">as</span> <span class="nn">sns</span>

        <span class="c1"># set seaborn style</span>
        <span class="n">sns</span><span class="o">.</span><span class="n">set_style</span><span class="p">(</span><span class="s2">&quot;white&quot;</span><span class="p">)</span>

        <span class="bp">self</span><span class="o">.</span><span class="n">matrix_lags</span> <span class="o">=</span> <span class="n">matrix_lags</span>

        <span class="k">if</span> <span class="n">matrix_lags</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> <span class="n">np</span><span class="o">.</span><span class="n">any</span><span class="p">(</span><span class="n">matrix_lags</span> <span class="o">&lt;</span> <span class="mi">0</span><span class="p">):</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;matrix_lags must be non-negative!&quot;</span><span class="p">)</span>

        <span class="n">data</span> <span class="o">=</span> <span class="n">dataframe</span><span class="o">.</span><span class="n">values</span><span class="p">[</span><span class="n">selected_dataset</span><span class="p">]</span>
        <span class="k">if</span> <span class="n">dataframe</span><span class="o">.</span><span class="n">mask</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">mask</span> <span class="o">=</span> <span class="n">dataframe</span><span class="o">.</span><span class="n">mask</span><span class="p">[</span><span class="n">selected_dataset</span><span class="p">]</span>

        <span class="n">T</span><span class="p">,</span> <span class="n">dim</span> <span class="o">=</span> <span class="n">data</span><span class="o">.</span><span class="n">shape</span>

        <span class="c1"># if label is not None:</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">labels</span><span class="o">.</span><span class="n">append</span><span class="p">((</span><span class="n">label</span><span class="p">,</span> <span class="n">label_color</span><span class="p">))</span>

        <span class="k">for</span> <span class="n">ij</span> <span class="ow">in</span> <span class="nb">list</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">):</span>  
            <span class="n">i</span> <span class="o">=</span> <span class="n">ij</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
            <span class="n">j</span> <span class="o">=</span> <span class="n">ij</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span>
            <span class="n">ax</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">[(</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">)]</span>              
            <span class="k">if</span> <span class="p">(</span><span class="n">matrix_lags</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">):</span>
                <span class="k">if</span> <span class="n">i</span> <span class="o">==</span> <span class="n">j</span><span class="p">:</span>
                    <span class="n">lag</span> <span class="o">=</span> <span class="mi">1</span>
                <span class="k">else</span><span class="p">:</span>
                    <span class="n">lag</span> <span class="o">=</span> <span class="mi">0</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">lag</span> <span class="o">=</span> <span class="n">matrix_lags</span><span class="p">[</span><span class="n">i</span><span class="p">,</span><span class="n">j</span><span class="p">]</span>
            <span class="n">x</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">copy</span><span class="p">(</span><span class="n">data</span><span class="p">[:</span><span class="n">T</span><span class="o">-</span><span class="n">lag</span><span class="p">,</span> <span class="n">i</span><span class="p">])</span>
            <span class="n">y</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">copy</span><span class="p">(</span><span class="n">data</span><span class="p">[</span><span class="n">lag</span><span class="p">:,</span> <span class="n">j</span><span class="p">])</span>
            <span class="c1"># Data is set to NaN in dataframe init already</span>
            <span class="c1"># if dataframe.missing_flag is not None:</span>
            <span class="c1">#     x[x==dataframe.missing_flag] = np.nan</span>
            <span class="c1">#     y[y==dataframe.missing_flag] = np.nan</span>
            <span class="k">if</span> <span class="n">dataframe</span><span class="o">.</span><span class="n">mask</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                <span class="n">x</span><span class="p">[</span><span class="n">mask</span><span class="p">[:</span><span class="n">T</span><span class="o">-</span><span class="n">lag</span><span class="p">,</span> <span class="n">i</span><span class="p">]</span><span class="o">==</span><span class="mi">1</span><span class="p">]</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">nan</span>
                <span class="n">y</span><span class="p">[</span><span class="n">mask</span><span class="p">[</span><span class="n">lag</span><span class="p">:,</span> <span class="n">j</span><span class="p">]</span><span class="o">==</span><span class="mi">1</span><span class="p">]</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">nan</span>

            <span class="k">if</span> <span class="n">i</span> <span class="o">==</span> <span class="n">j</span> <span class="ow">and</span> <span class="n">show_marginal_densities_on_diagonal</span><span class="p">:</span>
                <span class="n">sns</span><span class="o">.</span><span class="n">kdeplot</span><span class="p">(</span><span class="n">x</span><span class="p">,</span>
                    <span class="n">color</span> <span class="o">=</span> <span class="n">label_color</span><span class="p">,</span>
                    <span class="c1"># label=r&quot;$\tau{=}%d$&quot; %lag,</span>
                    <span class="o">**</span><span class="n">snskdeplot_diagonal_args</span><span class="p">,</span>
                    <span class="n">ax</span> <span class="o">=</span> <span class="n">ax</span><span class="p">)</span>
                <span class="n">ax</span><span class="o">.</span><span class="n">set_ylabel</span><span class="p">(</span><span class="s2">&quot;&quot;</span><span class="p">)</span>
                <span class="c1"># ax.yaxis.set_ticks_position(&quot;none&quot;)</span>
                <span class="c1"># ax.yaxis.set_ticklabels([])</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">sns</span><span class="o">.</span><span class="n">kdeplot</span><span class="p">(</span><span class="n">x</span><span class="o">=</span><span class="n">x</span><span class="p">,</span> <span class="n">y</span><span class="o">=</span><span class="n">y</span><span class="p">,</span> 
                    <span class="c1">#label=r&quot;$\tau{=}%d$&quot; %lag,</span>
                    <span class="o">**</span><span class="n">snskdeplot_args</span><span class="p">,</span>
                    <span class="c1"># fill=True,</span>
                    <span class="c1"># alpha=0.3,</span>
                    <span class="n">ax</span> <span class="o">=</span> <span class="n">ax</span><span class="p">)</span></div>

<div class="viewcode-block" id="setup_density_matrix.adjustfig"><a class="viewcode-back" href="../../index.html#tigramite.plotting.setup_density_matrix.adjustfig">[docs]</a>    <span class="k">def</span> <span class="nf">adjustfig</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">name</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">show_labels</span><span class="o">=</span><span class="kc">True</span><span class="p">):</span>
<span class="w">        </span><span class="sd">&quot;&quot;&quot;Adjust matrix figure.</span>

<span class="sd">        Parameters</span>
<span class="sd">        ----------</span>
<span class="sd">        name : str, optional (default: None)</span>
<span class="sd">            File name. If None, figure is shown in window.</span>
<span class="sd">        &quot;&quot;&quot;</span>

        <span class="c1"># Trick to plot legends</span>
        <span class="c1"># colors = []</span>
        <span class="c1"># for item in self.labels:</span>
        <span class="c1">#     colors.append(item[1])</span>
        <span class="k">for</span> <span class="n">ij</span> <span class="ow">in</span> <span class="nb">list</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">):</span>                
            <span class="n">i</span> <span class="o">=</span> <span class="n">ij</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
            <span class="n">j</span> <span class="o">=</span> <span class="n">ij</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span>
            <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">matrix_lags</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
                <span class="n">lag</span> <span class="o">=</span> <span class="mi">0</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">lag</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">matrix_lags</span><span class="p">[</span><span class="n">i</span><span class="p">,</span><span class="n">j</span><span class="p">]</span>
            <span class="k">if</span> <span class="n">i</span> <span class="o">!=</span> <span class="n">j</span><span class="p">:</span>
                <span class="n">colors</span> <span class="o">=</span> <span class="p">[]</span>
                <span class="k">for</span> <span class="n">item</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">labels</span><span class="p">:</span>
                    <span class="n">color</span> <span class="o">=</span> <span class="n">item</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span>
                    <span class="n">colors</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">color</span><span class="p">)</span>
                    <span class="k">if</span> <span class="n">show_labels</span><span class="p">:</span>
                        <span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">[(</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">)]</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span>
                            <span class="p">[],</span>
                            <span class="p">[],</span>
                            <span class="n">linestyle</span><span class="o">=</span><span class="s2">&quot;&quot;</span><span class="p">,</span>
                            <span class="n">color</span><span class="o">=</span><span class="n">color</span><span class="p">,</span>
                            <span class="n">label</span><span class="o">=</span><span class="sa">r</span><span class="s2">&quot;$\tau{=}</span><span class="si">%d</span><span class="s2">$&quot;</span> <span class="o">%</span><span class="n">lag</span><span class="p">,</span>
                    <span class="p">)</span>
                <span class="c1"># print(&#39;here&#39;)</span>
                <span class="n">leg</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">axes_dict</span><span class="p">[(</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">)]</span><span class="o">.</span><span class="n">legend</span><span class="p">(</span>
                    <span class="c1"># loc=&quot;best&quot;,</span>
                    <span class="n">ncol</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span>
                    <span class="c1"># bbox_to_anchor=(1.05, 0.0, 0.1, 1.0),</span>
                    <span class="c1"># borderaxespad=0,</span>
                    <span class="n">fontsize</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">legend_fontsize</span><span class="o">-</span><span class="mi">2</span><span class="p">,</span>
                    <span class="n">labelcolor</span><span class="o">=</span><span class="n">colors</span><span class="p">,</span>
                    <span class="p">)</span><span class="o">.</span><span class="n">draw_frame</span><span class="p">(</span><span class="kc">False</span><span class="p">)</span>
        
            <span class="c1"># if i == j:</span>
            <span class="c1">#     # self.axes_dict[(i, j)].axis[&quot;left&quot;].set_visible(False)</span>
            <span class="c1">#     _make_nice_axes(ax=self.axes_dict[(i, j)], where=[&quot;bottom&quot;], skip=(1, 1))</span>
            <span class="c1"># else:</span>
            <span class="c1"># _make_nice_axes(ax=self.axes_dict[(i, j)], where=[&quot;left&quot;, &quot;bottom&quot;], skip=(1, 1))</span>

        <span class="k">if</span> <span class="n">show_labels</span> <span class="ow">and</span> <span class="nb">len</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">labels</span><span class="p">)</span> <span class="o">&gt;</span> <span class="mi">1</span><span class="p">:</span>
            <span class="n">axlegend</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">fig</span><span class="o">.</span><span class="n">add_subplot</span><span class="p">(</span><span class="mi">111</span><span class="p">,</span> <span class="n">frameon</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span>
            <span class="n">axlegend</span><span class="o">.</span><span class="n">spines</span><span class="p">[</span><span class="s2">&quot;left&quot;</span><span class="p">]</span><span class="o">.</span><span class="n">set_color</span><span class="p">(</span><span class="s2">&quot;none&quot;</span><span class="p">)</span>
            <span class="n">axlegend</span><span class="o">.</span><span class="n">spines</span><span class="p">[</span><span class="s2">&quot;right&quot;</span><span class="p">]</span><span class="o">.</span><span class="n">set_color</span><span class="p">(</span><span class="s2">&quot;none&quot;</span><span class="p">)</span>
            <span class="n">axlegend</span><span class="o">.</span><span class="n">spines</span><span class="p">[</span><span class="s2">&quot;bottom&quot;</span><span class="p">]</span><span class="o">.</span><span class="n">set_color</span><span class="p">(</span><span class="s2">&quot;none&quot;</span><span class="p">)</span>
            <span class="n">axlegend</span><span class="o">.</span><span class="n">spines</span><span class="p">[</span><span class="s2">&quot;top&quot;</span><span class="p">]</span><span class="o">.</span><span class="n">set_color</span><span class="p">(</span><span class="s2">&quot;none&quot;</span><span class="p">)</span>
            <span class="n">axlegend</span><span class="o">.</span><span class="n">set_xticks</span><span class="p">([])</span>
            <span class="n">axlegend</span><span class="o">.</span><span class="n">set_yticks</span><span class="p">([])</span>

            <span class="c1"># self.labels.append((label, color, marker, markersize, alpha))</span>
            <span class="k">for</span> <span class="n">item</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">labels</span><span class="p">:</span>
                <span class="n">label</span> <span class="o">=</span> <span class="n">item</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
                <span class="n">color</span> <span class="o">=</span> <span class="n">item</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span>

                <span class="n">axlegend</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span>
                    <span class="p">[],</span>
                    <span class="p">[],</span>
                    <span class="n">linestyle</span><span class="o">=</span><span class="s2">&quot;-&quot;</span><span class="p">,</span>
                    <span class="n">color</span><span class="o">=</span><span class="n">color</span><span class="p">,</span>
                    <span class="n">label</span><span class="o">=</span><span class="n">label</span><span class="p">,</span>
                <span class="p">)</span>
            <span class="n">axlegend</span><span class="o">.</span><span class="n">legend</span><span class="p">(</span>
                <span class="n">loc</span><span class="o">=</span><span class="s2">&quot;upper left&quot;</span><span class="p">,</span>
                <span class="n">ncol</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span>
                <span class="n">bbox_to_anchor</span><span class="o">=</span><span class="p">(</span><span class="mf">1.05</span><span class="p">,</span> <span class="mf">0.0</span><span class="p">,</span> <span class="mf">0.1</span><span class="p">,</span> <span class="mf">1.0</span><span class="p">),</span>
                <span class="n">borderaxespad</span><span class="o">=</span><span class="mi">0</span><span class="p">,</span>
                <span class="n">fontsize</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">legend_fontsize</span><span class="p">,</span>
            <span class="p">)</span><span class="o">.</span><span class="n">draw_frame</span><span class="p">(</span><span class="kc">False</span><span class="p">)</span>

            <span class="bp">self</span><span class="o">.</span><span class="n">fig</span><span class="o">.</span><span class="n">subplots_adjust</span><span class="p">(</span>
                <span class="n">bottom</span><span class="o">=</span><span class="mf">0.08</span><span class="p">,</span>
                <span class="n">left</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">label_space_left</span><span class="p">,</span>
                <span class="n">right</span><span class="o">=</span><span class="mf">1.0</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">legend_width</span><span class="p">,</span>
                <span class="n">top</span><span class="o">=</span><span class="mf">1.0</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">label_space_top</span><span class="p">,</span>
                <span class="n">hspace</span><span class="o">=</span><span class="mf">0.5</span><span class="p">,</span>
                <span class="n">wspace</span><span class="o">=</span><span class="mf">0.35</span><span class="p">,</span>
            <span class="p">)</span>
      
        <span class="k">else</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">fig</span><span class="o">.</span><span class="n">subplots_adjust</span><span class="p">(</span>
                <span class="n">left</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">label_space_left</span><span class="p">,</span>
                <span class="n">bottom</span><span class="o">=</span><span class="mf">0.08</span><span class="p">,</span>
                <span class="n">right</span><span class="o">=</span><span class="mf">0.95</span><span class="p">,</span>
                <span class="n">top</span><span class="o">=</span><span class="mf">1.0</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">label_space_top</span><span class="p">,</span>
                <span class="n">hspace</span><span class="o">=</span><span class="mf">0.35</span><span class="p">,</span>
                <span class="n">wspace</span><span class="o">=</span><span class="mf">0.35</span><span class="p">,</span>
            <span class="p">)</span>
       
        <span class="k">if</span> <span class="n">name</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">fig</span><span class="o">.</span><span class="n">savefig</span><span class="p">(</span><span class="n">name</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">pyplot</span><span class="o">.</span><span class="n">show</span><span class="p">()</span></div></div>

<span class="k">def</span> <span class="nf">_draw_network_with_curved_edges</span><span class="p">(</span>
        <span class="n">fig</span><span class="p">,</span>
        <span class="n">ax</span><span class="p">,</span>
        <span class="n">G</span><span class="p">,</span>
        <span class="n">pos</span><span class="p">,</span>
        <span class="n">node_rings</span><span class="p">,</span>
        <span class="n">node_labels</span><span class="p">,</span>
        <span class="n">node_label_size</span><span class="o">=</span><span class="mi">10</span><span class="p">,</span>
        <span class="n">node_alpha</span><span class="o">=</span><span class="mf">1.0</span><span class="p">,</span>
        <span class="n">standard_size</span><span class="o">=</span><span class="mi">100</span><span class="p">,</span>
        <span class="n">node_aspect</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
        <span class="n">standard_cmap</span><span class="o">=</span><span class="s2">&quot;OrRd&quot;</span><span class="p">,</span>
        <span class="n">standard_color_links</span><span class="o">=</span><span class="s1">&#39;black&#39;</span><span class="p">,</span>
        <span class="n">standard_color_nodes</span><span class="o">=</span><span class="s1">&#39;lightgrey&#39;</span><span class="p">,</span>
        <span class="n">log_sizes</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
        <span class="n">cmap_links</span><span class="o">=</span><span class="s2">&quot;YlOrRd&quot;</span><span class="p">,</span>
        <span class="c1"># cmap_links_edges=&quot;YlOrRd&quot;,</span>
        <span class="n">links_vmin</span><span class="o">=</span><span class="mf">0.0</span><span class="p">,</span>
        <span class="n">links_vmax</span><span class="o">=</span><span class="mf">1.0</span><span class="p">,</span>
        <span class="n">links_edges_vmin</span><span class="o">=</span><span class="mf">0.0</span><span class="p">,</span>
        <span class="n">links_edges_vmax</span><span class="o">=</span><span class="mf">1.0</span><span class="p">,</span>
        <span class="n">links_ticks</span><span class="o">=</span><span class="mf">0.2</span><span class="p">,</span>
        <span class="n">links_edges_ticks</span><span class="o">=</span><span class="mf">0.2</span><span class="p">,</span>
        <span class="n">link_label_fontsize</span><span class="o">=</span><span class="mi">8</span><span class="p">,</span>
        <span class="n">arrowstyle</span><span class="o">=</span><span class="s2">&quot;-&gt;, head_width=0.4, head_length=1&quot;</span><span class="p">,</span>
        <span class="n">arrowhead_size</span><span class="o">=</span><span class="mf">3.0</span><span class="p">,</span>
        <span class="n">curved_radius</span><span class="o">=</span><span class="mf">0.2</span><span class="p">,</span>
        <span class="n">label_fontsize</span><span class="o">=</span><span class="mi">4</span><span class="p">,</span>
        <span class="n">label_fraction</span><span class="o">=</span><span class="mf">0.5</span><span class="p">,</span>
        <span class="n">link_colorbar_label</span><span class="o">=</span><span class="s2">&quot;link&quot;</span><span class="p">,</span>
        <span class="n">tick_label_size</span><span class="o">=</span><span class="mi">6</span><span class="p">,</span>
        <span class="c1"># link_edge_colorbar_label=&#39;link_edge&#39;,</span>
        <span class="n">inner_edge_curved</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
        <span class="n">inner_edge_style</span><span class="o">=</span><span class="s2">&quot;solid&quot;</span><span class="p">,</span>
        <span class="c1"># network_lower_bound=0.2,</span>
        <span class="n">network_left_bound</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
        <span class="n">show_colorbar</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
        <span class="n">special_nodes</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
        <span class="n">autodep_sig_lags</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
        <span class="n">show_autodependency_lags</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
        <span class="n">transform</span><span class="o">=</span><span class="s1">&#39;data&#39;</span><span class="p">,</span>
        <span class="n">node_classification</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
        <span class="n">max_lag</span><span class="o">=</span><span class="mi">0</span><span class="p">,</span>
<span class="p">):</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;Function to draw a network from networkx graph instance.</span>
<span class="sd">    Various attributes are used to specify the graph&#39;s properties.</span>
<span class="sd">    This function is just a beta-template for now that can be further</span>
<span class="sd">    customized.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">if</span> <span class="n">transform</span> <span class="o">==</span> <span class="s1">&#39;data&#39;</span><span class="p">:</span>
        <span class="n">transform</span> <span class="o">=</span> <span class="n">ax</span><span class="o">.</span><span class="n">transData</span>

    <span class="kn">from</span> <span class="nn">matplotlib.patches</span> <span class="kn">import</span> <span class="n">FancyArrowPatch</span><span class="p">,</span> <span class="n">Circle</span><span class="p">,</span> <span class="n">Ellipse</span>

    <span class="n">ax</span><span class="o">.</span><span class="n">spines</span><span class="p">[</span><span class="s2">&quot;left&quot;</span><span class="p">]</span><span class="o">.</span><span class="n">set_color</span><span class="p">(</span><span class="s2">&quot;none&quot;</span><span class="p">)</span>
    <span class="n">ax</span><span class="o">.</span><span class="n">spines</span><span class="p">[</span><span class="s2">&quot;right&quot;</span><span class="p">]</span><span class="o">.</span><span class="n">set_color</span><span class="p">(</span><span class="s2">&quot;none&quot;</span><span class="p">)</span>
    <span class="n">ax</span><span class="o">.</span><span class="n">spines</span><span class="p">[</span><span class="s2">&quot;bottom&quot;</span><span class="p">]</span><span class="o">.</span><span class="n">set_color</span><span class="p">(</span><span class="s2">&quot;none&quot;</span><span class="p">)</span>
    <span class="n">ax</span><span class="o">.</span><span class="n">spines</span><span class="p">[</span><span class="s2">&quot;top&quot;</span><span class="p">]</span><span class="o">.</span><span class="n">set_color</span><span class="p">(</span><span class="s2">&quot;none&quot;</span><span class="p">)</span>
    <span class="n">ax</span><span class="o">.</span><span class="n">set_xticks</span><span class="p">([])</span>
    <span class="n">ax</span><span class="o">.</span><span class="n">set_yticks</span><span class="p">([])</span>

    <span class="n">N</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="n">G</span><span class="p">)</span>

    <span class="c1"># This fixes a positioning bug in matplotlib.</span>
    <span class="n">ax</span><span class="o">.</span><span class="n">scatter</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="n">zorder</span><span class="o">=-</span><span class="mi">10</span><span class="p">,</span> <span class="n">alpha</span><span class="o">=</span><span class="mi">0</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">draw_edge</span><span class="p">(</span>
        <span class="n">ax</span><span class="p">,</span>
        <span class="n">u</span><span class="p">,</span>
        <span class="n">v</span><span class="p">,</span>
        <span class="n">d</span><span class="p">,</span>
        <span class="n">seen</span><span class="p">,</span>
        <span class="n">arrowstyle</span><span class="o">=</span> <span class="s2">&quot;Simple, head_width=2, head_length=2, tail_width=1&quot;</span><span class="p">,</span>
        <span class="n">outer_edge</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
    <span class="p">):</span>

        <span class="c1"># avoiding attribute error raised by changes in networkx</span>
        <span class="k">if</span> <span class="nb">hasattr</span><span class="p">(</span><span class="n">G</span><span class="p">,</span> <span class="s2">&quot;node&quot;</span><span class="p">):</span>
            <span class="c1"># works with networkx 1.10</span>
            <span class="n">n1</span> <span class="o">=</span> <span class="n">G</span><span class="o">.</span><span class="n">node</span><span class="p">[</span><span class="n">u</span><span class="p">][</span><span class="s2">&quot;patch&quot;</span><span class="p">]</span>
            <span class="n">n2</span> <span class="o">=</span> <span class="n">G</span><span class="o">.</span><span class="n">node</span><span class="p">[</span><span class="n">v</span><span class="p">][</span><span class="s2">&quot;patch&quot;</span><span class="p">]</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="c1"># works with networkx 2.4</span>
            <span class="n">n1</span> <span class="o">=</span> <span class="n">G</span><span class="o">.</span><span class="n">nodes</span><span class="p">[</span><span class="n">u</span><span class="p">][</span><span class="s2">&quot;patch&quot;</span><span class="p">]</span>
            <span class="n">n2</span> <span class="o">=</span> <span class="n">G</span><span class="o">.</span><span class="n">nodes</span><span class="p">[</span><span class="n">v</span><span class="p">][</span><span class="s2">&quot;patch&quot;</span><span class="p">]</span>

        <span class="c1"># print(&quot;+++++++++++++++++++++++==cmap_links &quot;, cmap_links)</span>
        <span class="k">if</span> <span class="n">outer_edge</span><span class="p">:</span>
            <span class="n">rad</span> <span class="o">=</span> <span class="o">-</span><span class="mf">1.0</span> <span class="o">*</span> <span class="n">curved_radius</span>
            <span class="k">if</span> <span class="n">cmap_links</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                <span class="n">facecolor</span> <span class="o">=</span> <span class="n">data_to_rgb_links</span><span class="o">.</span><span class="n">to_rgba</span><span class="p">(</span><span class="n">d</span><span class="p">[</span><span class="s2">&quot;outer_edge_color&quot;</span><span class="p">])</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="k">if</span> <span class="n">d</span><span class="p">[</span><span class="s2">&quot;outer_edge_color&quot;</span><span class="p">]</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                    <span class="n">facecolor</span> <span class="o">=</span> <span class="n">d</span><span class="p">[</span><span class="s2">&quot;outer_edge_color&quot;</span><span class="p">]</span>
                <span class="k">else</span><span class="p">:</span>
                    <span class="n">facecolor</span> <span class="o">=</span> <span class="n">standard_color_links</span>

            <span class="n">width</span> <span class="o">=</span> <span class="n">d</span><span class="p">[</span><span class="s2">&quot;outer_edge_width&quot;</span><span class="p">]</span>
            <span class="n">alpha</span> <span class="o">=</span> <span class="n">d</span><span class="p">[</span><span class="s2">&quot;outer_edge_alpha&quot;</span><span class="p">]</span>
            <span class="k">if</span> <span class="p">(</span><span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">)</span> <span class="ow">in</span> <span class="n">seen</span><span class="p">:</span>
                <span class="n">rad</span> <span class="o">=</span> <span class="n">seen</span><span class="o">.</span><span class="n">get</span><span class="p">((</span><span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">))</span>
                <span class="n">rad</span> <span class="o">=</span> <span class="p">(</span><span class="n">rad</span> <span class="o">+</span> <span class="n">np</span><span class="o">.</span><span class="n">sign</span><span class="p">(</span><span class="n">rad</span><span class="p">)</span> <span class="o">*</span> <span class="mf">0.1</span><span class="p">)</span> <span class="o">*</span> <span class="o">-</span><span class="mf">1.0</span>
            <span class="n">arrowstyle</span> <span class="o">=</span> <span class="n">arrowstyle</span>
            <span class="c1"># link_edge = d[&#39;outer_edge_edge&#39;]</span>
            <span class="n">linestyle</span> <span class="o">=</span> <span class="s1">&#39;solid&#39;</span> <span class="c1"># d.get(&quot;outer_edge_style&quot;)</span>

            <span class="k">if</span> <span class="n">d</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;outer_edge_attribute&quot;</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span> <span class="o">==</span> <span class="s2">&quot;spurious&quot;</span><span class="p">:</span>
                <span class="n">facecolor</span> <span class="o">=</span> <span class="s2">&quot;grey&quot;</span>

            <span class="k">if</span> <span class="n">d</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;outer_edge_type&quot;</span><span class="p">)</span> <span class="ow">in</span> <span class="p">[</span><span class="s2">&quot;&lt;-o&quot;</span><span class="p">,</span> <span class="s2">&quot;&lt;--&quot;</span><span class="p">,</span> <span class="s2">&quot;&lt;-x&quot;</span><span class="p">,</span> <span class="s2">&quot;&lt;-+&quot;</span><span class="p">]:</span>
                <span class="n">n1</span><span class="p">,</span> <span class="n">n2</span> <span class="o">=</span> <span class="n">n2</span><span class="p">,</span> <span class="n">n1</span>

            <span class="k">if</span> <span class="n">d</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;outer_edge_type&quot;</span><span class="p">)</span> <span class="ow">in</span> <span class="p">[</span>
                <span class="s2">&quot;o-o&quot;</span><span class="p">,</span>
                <span class="s2">&quot;o--&quot;</span><span class="p">,</span>
                <span class="s2">&quot;--o&quot;</span><span class="p">,</span>
                <span class="s2">&quot;---&quot;</span><span class="p">,</span>
                <span class="s2">&quot;x-x&quot;</span><span class="p">,</span>
                <span class="s2">&quot;x--&quot;</span><span class="p">,</span>
                <span class="s2">&quot;--x&quot;</span><span class="p">,</span>
                <span class="s2">&quot;o-x&quot;</span><span class="p">,</span>
                <span class="s2">&quot;x-o&quot;</span><span class="p">,</span>
                <span class="c1"># &quot;+-&gt;&quot;,</span>
                <span class="c1"># &quot;&lt;-+&quot;,</span>
            <span class="p">]:</span>
                <span class="n">arrowstyle</span> <span class="o">=</span> <span class="s2">&quot;-&quot;</span>
                <span class="c1"># linewidth = width*factor</span>
            <span class="k">elif</span> <span class="n">d</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;outer_edge_type&quot;</span><span class="p">)</span> <span class="o">==</span> <span class="s2">&quot;&lt;-&gt;&quot;</span><span class="p">:</span>
                <span class="c1"># arrowstyle = &quot;&lt;-&gt;, head_width=0.4, head_length=1&quot;</span>
                <span class="n">arrowstyle</span> <span class="o">=</span> <span class="s2">&quot;Simple, head_width=2, head_length=2, tail_width=1&quot;</span> <span class="c1">#%float(width/20.)</span>
            <span class="k">elif</span> <span class="n">d</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;outer_edge_type&quot;</span><span class="p">)</span> <span class="ow">in</span> <span class="p">[</span><span class="s2">&quot;o-&gt;&quot;</span><span class="p">,</span> <span class="s2">&quot;--&gt;&quot;</span><span class="p">,</span> <span class="s2">&quot;&lt;-o&quot;</span><span class="p">,</span> <span class="s2">&quot;&lt;--&quot;</span><span class="p">,</span> <span class="s2">&quot;&lt;-x&quot;</span><span class="p">,</span> <span class="s2">&quot;x-&gt;&quot;</span><span class="p">,</span> <span class="s2">&quot;+-&gt;&quot;</span><span class="p">,</span> <span class="s2">&quot;&lt;-+&quot;</span><span class="p">]:</span>
                <span class="c1"># arrowstyle = &quot;-&gt;, head_width=0.4, head_length=1&quot;</span>
                <span class="c1"># arrowstyle = &quot;-&gt;, head_width=0.4, head_length=1, width=10&quot;</span>
                <span class="n">arrowstyle</span> <span class="o">=</span> <span class="s2">&quot;Simple, head_width=2, head_length=2, tail_width=1&quot;</span> <span class="c1">#%float(width/20.)</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">arrowstyle</span> <span class="o">=</span> <span class="s2">&quot;Simple, head_width=2, head_length=2, tail_width=1&quot;</span> <span class="c1">#%float(width/20.)</span>
                <span class="c1"># raise ValueError(&quot;edge type %s not valid.&quot; %d.get(&quot;outer_edge_type&quot;))</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">rad</span> <span class="o">=</span> <span class="o">-</span><span class="mf">1.0</span> <span class="o">*</span> <span class="n">inner_edge_curved</span> <span class="o">*</span> <span class="n">curved_radius</span>
            <span class="k">if</span> <span class="n">cmap_links</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                <span class="n">facecolor</span> <span class="o">=</span> <span class="n">data_to_rgb_links</span><span class="o">.</span><span class="n">to_rgba</span><span class="p">(</span><span class="n">d</span><span class="p">[</span><span class="s2">&quot;inner_edge_color&quot;</span><span class="p">])</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="k">if</span> <span class="n">d</span><span class="p">[</span><span class="s2">&quot;inner_edge_color&quot;</span><span class="p">]</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                    <span class="n">facecolor</span> <span class="o">=</span> <span class="n">d</span><span class="p">[</span><span class="s2">&quot;inner_edge_color&quot;</span><span class="p">]</span>
                <span class="k">else</span><span class="p">:</span>
                    <span class="c1"># print(&quot;HERE&quot;)</span>
                    <span class="n">facecolor</span> <span class="o">=</span> <span class="n">standard_color_links</span>

            <span class="n">width</span> <span class="o">=</span> <span class="n">d</span><span class="p">[</span><span class="s2">&quot;inner_edge_width&quot;</span><span class="p">]</span>
            <span class="n">alpha</span> <span class="o">=</span> <span class="n">d</span><span class="p">[</span><span class="s2">&quot;inner_edge_alpha&quot;</span><span class="p">]</span>

            <span class="k">if</span> <span class="n">d</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;inner_edge_attribute&quot;</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span> <span class="o">==</span> <span class="s2">&quot;spurious&quot;</span><span class="p">:</span>
                <span class="n">facecolor</span> <span class="o">=</span> <span class="s2">&quot;grey&quot;</span>
            <span class="c1"># print(d.get(&quot;inner_edge_type&quot;))</span>
            <span class="k">if</span> <span class="n">d</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;inner_edge_type&quot;</span><span class="p">)</span> <span class="ow">in</span> <span class="p">[</span><span class="s2">&quot;&lt;-o&quot;</span><span class="p">,</span> <span class="s2">&quot;&lt;--&quot;</span><span class="p">,</span> <span class="s2">&quot;&lt;-x&quot;</span><span class="p">,</span> <span class="s2">&quot;&lt;-+&quot;</span><span class="p">]:</span>
                <span class="n">n1</span><span class="p">,</span> <span class="n">n2</span> <span class="o">=</span> <span class="n">n2</span><span class="p">,</span> <span class="n">n1</span>

            <span class="k">if</span> <span class="n">d</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;inner_edge_type&quot;</span><span class="p">)</span> <span class="ow">in</span> <span class="p">[</span>
                <span class="s2">&quot;o-o&quot;</span><span class="p">,</span>
                <span class="s2">&quot;o--&quot;</span><span class="p">,</span>
                <span class="s2">&quot;--o&quot;</span><span class="p">,</span>
                <span class="s2">&quot;---&quot;</span><span class="p">,</span>
                <span class="s2">&quot;x-x&quot;</span><span class="p">,</span>
                <span class="s2">&quot;x--&quot;</span><span class="p">,</span>
                <span class="s2">&quot;--x&quot;</span><span class="p">,</span>
                <span class="s2">&quot;o-x&quot;</span><span class="p">,</span>
                <span class="s2">&quot;x-o&quot;</span><span class="p">,</span>
            <span class="p">]:</span>
                <span class="n">arrowstyle</span> <span class="o">=</span> <span class="s2">&quot;-&quot;</span>
            <span class="k">elif</span> <span class="n">d</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;inner_edge_type&quot;</span><span class="p">)</span> <span class="o">==</span> <span class="s2">&quot;&lt;-&gt;&quot;</span><span class="p">:</span>
                <span class="c1"># arrowstyle = &quot;&lt;-&gt;, head_width=0.4, head_length=1&quot;</span>
                <span class="n">arrowstyle</span> <span class="o">=</span> <span class="s2">&quot;Simple, head_width=2, head_length=2, tail_width=1&quot;</span> <span class="c1">#%float(width/20.)</span>
            <span class="k">elif</span> <span class="n">d</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;inner_edge_type&quot;</span><span class="p">)</span> <span class="ow">in</span> <span class="p">[</span><span class="s2">&quot;o-&gt;&quot;</span><span class="p">,</span> <span class="s2">&quot;--&gt;&quot;</span><span class="p">,</span> <span class="s2">&quot;&lt;-o&quot;</span><span class="p">,</span> <span class="s2">&quot;&lt;--&quot;</span><span class="p">,</span> <span class="s2">&quot;&lt;-x&quot;</span><span class="p">,</span> <span class="s2">&quot;x-&gt;&quot;</span><span class="p">,</span> <span class="s2">&quot;+-&gt;&quot;</span><span class="p">,</span> <span class="s2">&quot;&lt;-+&quot;</span><span class="p">]:</span>
                <span class="c1"># arrowstyle = &quot;-&gt;, head_width=0.4, head_length=1&quot;</span>
                <span class="n">arrowstyle</span> <span class="o">=</span> <span class="s2">&quot;Simple, head_width=2, head_length=2, tail_width=1&quot;</span> <span class="c1">#%float(width/20.)</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">arrowstyle</span> <span class="o">=</span> <span class="s2">&quot;Simple, head_width=2, head_length=2, tail_width=1&quot;</span> <span class="c1">#%float(width/20.)</span>

            <span class="c1">#     raise ValueError(&quot;edge type %s not valid.&quot; %d.get(&quot;inner_edge_type&quot;))</span>

            <span class="n">linestyle</span> <span class="o">=</span> <span class="s1">&#39;solid&#39;</span> <span class="c1">#d.get(&quot;inner_edge_style&quot;)</span>

        <span class="n">coor1</span> <span class="o">=</span> <span class="n">n1</span><span class="o">.</span><span class="n">center</span>
        <span class="n">coor2</span> <span class="o">=</span> <span class="n">n2</span><span class="o">.</span><span class="n">center</span>

        <span class="n">marker_size</span> <span class="o">=</span> <span class="n">width</span> <span class="o">**</span> <span class="mi">2</span>
        <span class="n">figuresize</span> <span class="o">=</span> <span class="n">fig</span><span class="o">.</span><span class="n">get_size_inches</span><span class="p">()</span>

        <span class="c1"># print(&quot;COLOR &quot;, facecolor)</span>
        <span class="c1"># print(u, v, outer_edge, &quot;outer &quot;, d.get(&quot;outer_edge_type&quot;),  &quot;inner &quot;,  d.get(&quot;inner_edge_type&quot;), width, arrowstyle, linestyle)</span>
        
        <span class="k">if</span> <span class="p">((</span><span class="n">outer_edge</span> <span class="ow">is</span> <span class="kc">True</span> <span class="ow">and</span> <span class="n">d</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;outer_edge_type&quot;</span><span class="p">)</span> <span class="o">==</span> <span class="s2">&quot;&lt;-&gt;&quot;</span><span class="p">)</span>
           <span class="ow">or</span> <span class="p">(</span><span class="n">outer_edge</span> <span class="ow">is</span> <span class="kc">False</span> <span class="ow">and</span> <span class="n">d</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;inner_edge_type&quot;</span><span class="p">)</span> <span class="o">==</span> <span class="s2">&quot;&lt;-&gt;&quot;</span><span class="p">)):</span>
            <span class="n">e_p</span> <span class="o">=</span> <span class="n">FancyArrowPatch</span><span class="p">(</span>
                <span class="n">coor1</span><span class="p">,</span>
                <span class="n">coor2</span><span class="p">,</span>
                <span class="n">arrowstyle</span><span class="o">=</span><span class="n">arrowstyle</span><span class="p">,</span>
                <span class="n">connectionstyle</span><span class="o">=</span><span class="sa">f</span><span class="s2">&quot;arc3,rad=</span><span class="si">{</span><span class="n">rad</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">,</span>
                <span class="n">mutation_scale</span><span class="o">=</span><span class="mi">1</span><span class="o">*</span><span class="n">width</span><span class="p">,</span>
                <span class="n">lw</span><span class="o">=</span><span class="mf">0.</span><span class="p">,</span> <span class="c1">#width / 2.,</span>
                <span class="n">aa</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
                <span class="n">alpha</span><span class="o">=</span><span class="n">alpha</span><span class="p">,</span>
                <span class="n">linestyle</span><span class="o">=</span><span class="n">linestyle</span><span class="p">,</span>
                <span class="n">color</span><span class="o">=</span><span class="n">facecolor</span><span class="p">,</span>
                <span class="n">clip_on</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
                <span class="n">patchA</span><span class="o">=</span><span class="n">n1</span><span class="p">,</span>
                <span class="n">patchB</span><span class="o">=</span><span class="n">n2</span><span class="p">,</span>
                <span class="n">shrinkA</span><span class="o">=</span><span class="mi">7</span><span class="p">,</span>
                <span class="n">shrinkB</span><span class="o">=</span><span class="mi">0</span><span class="p">,</span>
                <span class="n">zorder</span><span class="o">=-</span><span class="mi">1</span><span class="p">,</span>
                <span class="n">capstyle</span><span class="o">=</span><span class="s2">&quot;butt&quot;</span><span class="p">,</span>
                <span class="n">transform</span><span class="o">=</span><span class="n">transform</span><span class="p">,</span>
            <span class="p">)</span>
            <span class="n">ax</span><span class="o">.</span><span class="n">add_artist</span><span class="p">(</span><span class="n">e_p</span><span class="p">)</span>

            <span class="n">e_p_back</span> <span class="o">=</span> <span class="n">FancyArrowPatch</span><span class="p">(</span>
              <span class="n">coor2</span><span class="p">,</span>
              <span class="n">coor1</span><span class="p">,</span>
              <span class="n">arrowstyle</span><span class="o">=</span><span class="n">arrowstyle</span><span class="p">,</span>
              <span class="n">connectionstyle</span><span class="o">=</span><span class="sa">f</span><span class="s2">&quot;arc3,rad=</span><span class="si">{</span><span class="o">-</span><span class="n">rad</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">,</span>
              <span class="n">mutation_scale</span><span class="o">=</span><span class="mi">1</span><span class="o">*</span><span class="n">width</span><span class="p">,</span>
              <span class="n">lw</span><span class="o">=</span><span class="mf">0.</span><span class="p">,</span> <span class="c1">#width / 2.,</span>
              <span class="n">aa</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
              <span class="n">alpha</span><span class="o">=</span><span class="n">alpha</span><span class="p">,</span>
              <span class="n">linestyle</span><span class="o">=</span><span class="n">linestyle</span><span class="p">,</span>
              <span class="n">color</span><span class="o">=</span><span class="n">facecolor</span><span class="p">,</span>
              <span class="n">clip_on</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
              <span class="n">patchA</span><span class="o">=</span><span class="n">n2</span><span class="p">,</span>
              <span class="n">patchB</span><span class="o">=</span><span class="n">n1</span><span class="p">,</span>
              <span class="n">shrinkA</span><span class="o">=</span><span class="mi">7</span><span class="p">,</span>
              <span class="n">shrinkB</span><span class="o">=</span><span class="mi">0</span><span class="p">,</span>
              <span class="n">zorder</span><span class="o">=-</span><span class="mi">1</span><span class="p">,</span>
              <span class="n">capstyle</span><span class="o">=</span><span class="s2">&quot;butt&quot;</span><span class="p">,</span>
              <span class="n">transform</span><span class="o">=</span><span class="n">transform</span><span class="p">,</span>
            <span class="p">)</span>  
            <span class="n">ax</span><span class="o">.</span><span class="n">add_artist</span><span class="p">(</span><span class="n">e_p_back</span><span class="p">)</span>

        <span class="k">else</span><span class="p">:</span>
            <span class="k">if</span> <span class="n">arrowstyle</span> <span class="o">==</span> <span class="s1">&#39;-&#39;</span><span class="p">:</span>
                <span class="n">lw</span> <span class="o">=</span> <span class="mi">1</span><span class="o">*</span><span class="n">width</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">lw</span> <span class="o">=</span> <span class="mf">0.</span>
            <span class="c1"># e_p = FancyArrowPatch(</span>
            <span class="c1">#     coor1,</span>
            <span class="c1">#     coor2,</span>
            <span class="c1">#     arrowstyle=arrowstyle,</span>
            <span class="c1">#     connectionstyle=f&quot;arc3,rad={rad}&quot;,</span>
            <span class="c1">#     mutation_scale=np.sqrt(width)*2*1.1,</span>
            <span class="c1">#     lw=lw*1.1, #width / 2.,</span>
            <span class="c1">#     aa=True,</span>
            <span class="c1">#     alpha=alpha,</span>
            <span class="c1">#     linestyle=linestyle,</span>
            <span class="c1">#     color=&#39;white&#39;,</span>
            <span class="c1">#     clip_on=False,</span>
            <span class="c1">#     patchA=n1,</span>
            <span class="c1">#     patchB=n2,</span>
            <span class="c1">#     shrinkA=0,</span>
            <span class="c1">#     shrinkB=0,</span>
            <span class="c1">#     zorder=-1,</span>
            <span class="c1">#     capstyle=&quot;butt&quot;,</span>
            <span class="c1"># )</span>
            <span class="c1"># ax.add_artist(e_p)</span>
            <span class="n">e_p</span> <span class="o">=</span> <span class="n">FancyArrowPatch</span><span class="p">(</span>
                <span class="n">coor1</span><span class="p">,</span>
                <span class="n">coor2</span><span class="p">,</span>
                <span class="n">arrowstyle</span><span class="o">=</span><span class="n">arrowstyle</span><span class="p">,</span>
                <span class="n">connectionstyle</span><span class="o">=</span><span class="sa">f</span><span class="s2">&quot;arc3,rad=</span><span class="si">{</span><span class="n">rad</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">,</span>
                <span class="n">mutation_scale</span><span class="o">=</span><span class="mi">1</span><span class="o">*</span><span class="n">width</span><span class="p">,</span>
                <span class="n">lw</span><span class="o">=</span><span class="n">lw</span><span class="p">,</span> <span class="c1">#width / 2.,</span>
                <span class="n">aa</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
                <span class="n">alpha</span><span class="o">=</span><span class="n">alpha</span><span class="p">,</span>
                <span class="n">linestyle</span><span class="o">=</span><span class="n">linestyle</span><span class="p">,</span>
                <span class="n">color</span><span class="o">=</span><span class="n">facecolor</span><span class="p">,</span>
                <span class="n">clip_on</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
                <span class="n">patchA</span><span class="o">=</span><span class="n">n1</span><span class="p">,</span>
                <span class="n">patchB</span><span class="o">=</span><span class="n">n2</span><span class="p">,</span>
                <span class="n">shrinkA</span><span class="o">=</span><span class="mi">0</span><span class="p">,</span>
                <span class="n">shrinkB</span><span class="o">=</span><span class="mi">0</span><span class="p">,</span>
                <span class="c1"># zorder=-1,</span>
                <span class="n">capstyle</span><span class="o">=</span><span class="s2">&quot;butt&quot;</span><span class="p">,</span>
                <span class="n">transform</span><span class="o">=</span><span class="n">transform</span><span class="p">,</span>
            <span class="p">)</span>
            <span class="n">ax</span><span class="o">.</span><span class="n">add_artist</span><span class="p">(</span><span class="n">e_p</span><span class="p">)</span>

        <span class="n">e_p_marker</span> <span class="o">=</span> <span class="n">FancyArrowPatch</span><span class="p">(</span>
                <span class="n">coor1</span><span class="p">,</span>
                <span class="n">coor2</span><span class="p">,</span>
                <span class="n">arrowstyle</span><span class="o">=</span><span class="s1">&#39;-&#39;</span><span class="p">,</span>
                <span class="n">connectionstyle</span><span class="o">=</span><span class="sa">f</span><span class="s2">&quot;arc3,rad=</span><span class="si">{</span><span class="n">rad</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">,</span>
                <span class="n">mutation_scale</span><span class="o">=</span><span class="mi">1</span><span class="o">*</span><span class="n">width</span><span class="p">,</span>
                <span class="n">lw</span><span class="o">=</span><span class="mf">0.</span><span class="p">,</span> <span class="c1">#width / 2.,</span>
                <span class="n">aa</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
                <span class="n">alpha</span><span class="o">=</span><span class="mf">0.</span><span class="p">,</span>
                <span class="n">linestyle</span><span class="o">=</span><span class="n">linestyle</span><span class="p">,</span>
                <span class="n">color</span><span class="o">=</span><span class="n">facecolor</span><span class="p">,</span>
                <span class="n">clip_on</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
                <span class="n">patchA</span><span class="o">=</span><span class="n">n1</span><span class="p">,</span>
                <span class="n">patchB</span><span class="o">=</span><span class="n">n2</span><span class="p">,</span>
                <span class="n">shrinkA</span><span class="o">=</span><span class="mi">0</span><span class="p">,</span>
                <span class="n">shrinkB</span><span class="o">=</span><span class="mi">0</span><span class="p">,</span>
                <span class="n">zorder</span><span class="o">=-</span><span class="mi">10</span><span class="p">,</span>
                <span class="n">capstyle</span><span class="o">=</span><span class="s2">&quot;butt&quot;</span><span class="p">,</span>
                <span class="n">transform</span><span class="o">=</span><span class="n">transform</span><span class="p">,</span>
        <span class="p">)</span>
        <span class="n">ax</span><span class="o">.</span><span class="n">add_artist</span><span class="p">(</span><span class="n">e_p_marker</span><span class="p">)</span>

        <span class="c1"># marker_path = e_p_marker.get_path()</span>
        <span class="n">vertices</span> <span class="o">=</span> <span class="n">e_p_marker</span><span class="o">.</span><span class="n">get_path</span><span class="p">()</span><span class="o">.</span><span class="n">vertices</span><span class="o">.</span><span class="n">copy</span><span class="p">()</span>
        <span class="c1"># vertices = e_p_marker.get_verts()</span>
        <span class="c1"># vertices = e_p_marker.get_path().to_polygons(transform=None)[0]</span>
        <span class="c1"># print(vertices.shape)</span>
        <span class="n">m</span><span class="p">,</span> <span class="n">n</span> <span class="o">=</span> <span class="n">vertices</span><span class="o">.</span><span class="n">shape</span>

        <span class="c1"># print(vertices)</span>
        <span class="n">start</span> <span class="o">=</span> <span class="n">vertices</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
        <span class="n">end</span> <span class="o">=</span> <span class="n">vertices</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span>

        <span class="c1"># This must be added to avoid rescaling of the plot, when no &#39;o&#39;</span>
        <span class="c1"># or &#39;x&#39; is added to the graph.</span>
        <span class="n">ax</span><span class="o">.</span><span class="n">scatter</span><span class="p">(</span><span class="o">*</span><span class="n">start</span><span class="p">,</span> <span class="n">zorder</span><span class="o">=-</span><span class="mi">10</span><span class="p">,</span> <span class="n">alpha</span><span class="o">=</span><span class="mi">0</span><span class="p">,</span> <span class="n">transform</span><span class="o">=</span><span class="n">transform</span><span class="p">,)</span>

        <span class="k">if</span> <span class="n">outer_edge</span><span class="p">:</span>
            <span class="k">if</span> <span class="n">d</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;outer_edge_type&quot;</span><span class="p">)</span> <span class="ow">in</span> <span class="p">[</span><span class="s2">&quot;o-&gt;&quot;</span><span class="p">,</span> <span class="s2">&quot;o--&quot;</span><span class="p">]:</span>
                <span class="n">circle_marker_start</span> <span class="o">=</span> <span class="n">ax</span><span class="o">.</span><span class="n">scatter</span><span class="p">(</span>
                    <span class="o">*</span><span class="n">start</span><span class="p">,</span>
                    <span class="n">marker</span><span class="o">=</span><span class="s2">&quot;o&quot;</span><span class="p">,</span>
                    <span class="n">s</span><span class="o">=</span><span class="n">marker_size</span><span class="p">,</span>
                    <span class="n">facecolor</span><span class="o">=</span><span class="s2">&quot;w&quot;</span><span class="p">,</span>
                    <span class="n">edgecolor</span><span class="o">=</span><span class="n">facecolor</span><span class="p">,</span>
                    <span class="n">zorder</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span>
                    <span class="n">transform</span><span class="o">=</span><span class="n">transform</span><span class="p">,</span>
                <span class="p">)</span>
                <span class="n">ax</span><span class="o">.</span><span class="n">add_collection</span><span class="p">(</span><span class="n">circle_marker_start</span><span class="p">)</span>
            <span class="k">elif</span> <span class="n">d</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;outer_edge_type&quot;</span><span class="p">)</span> <span class="o">==</span> <span class="s2">&quot;&lt;-o&quot;</span><span class="p">:</span>
                <span class="n">circle_marker_end</span> <span class="o">=</span> <span class="n">ax</span><span class="o">.</span><span class="n">scatter</span><span class="p">(</span>
                    <span class="o">*</span><span class="n">start</span><span class="p">,</span>
                    <span class="n">marker</span><span class="o">=</span><span class="s2">&quot;o&quot;</span><span class="p">,</span>
                    <span class="n">s</span><span class="o">=</span><span class="n">marker_size</span><span class="p">,</span>
                    <span class="n">facecolor</span><span class="o">=</span><span class="s2">&quot;w&quot;</span><span class="p">,</span>
                    <span class="n">edgecolor</span><span class="o">=</span><span class="n">facecolor</span><span class="p">,</span>
                    <span class="n">zorder</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span>
                    <span class="n">transform</span><span class="o">=</span><span class="n">transform</span><span class="p">,</span>
                <span class="p">)</span>
                <span class="n">ax</span><span class="o">.</span><span class="n">add_collection</span><span class="p">(</span><span class="n">circle_marker_end</span><span class="p">)</span>
            <span class="k">elif</span> <span class="n">d</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;outer_edge_type&quot;</span><span class="p">)</span> <span class="o">==</span> <span class="s2">&quot;--o&quot;</span><span class="p">:</span>
                <span class="n">circle_marker_end</span> <span class="o">=</span> <span class="n">ax</span><span class="o">.</span><span class="n">scatter</span><span class="p">(</span>
                    <span class="o">*</span><span class="n">end</span><span class="p">,</span>
                    <span class="n">marker</span><span class="o">=</span><span class="s2">&quot;o&quot;</span><span class="p">,</span>
                    <span class="n">s</span><span class="o">=</span><span class="n">marker_size</span><span class="p">,</span>
                    <span class="n">facecolor</span><span class="o">=</span><span class="s2">&quot;w&quot;</span><span class="p">,</span>
                    <span class="n">edgecolor</span><span class="o">=</span><span class="n">facecolor</span><span class="p">,</span>
                    <span class="n">zorder</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span>
                    <span class="n">transform</span><span class="o">=</span><span class="n">transform</span><span class="p">,</span>
                <span class="p">)</span>
                <span class="n">ax</span><span class="o">.</span><span class="n">add_collection</span><span class="p">(</span><span class="n">circle_marker_end</span><span class="p">)</span>
            <span class="k">elif</span> <span class="n">d</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;outer_edge_type&quot;</span><span class="p">)</span> <span class="ow">in</span> <span class="p">[</span><span class="s2">&quot;x--&quot;</span><span class="p">,</span> <span class="s2">&quot;x-&gt;&quot;</span><span class="p">]:</span>
                <span class="n">circle_marker_start</span> <span class="o">=</span> <span class="n">ax</span><span class="o">.</span><span class="n">scatter</span><span class="p">(</span>
                    <span class="o">*</span><span class="n">start</span><span class="p">,</span>
                    <span class="n">marker</span><span class="o">=</span><span class="s2">&quot;X&quot;</span><span class="p">,</span>
                    <span class="n">s</span><span class="o">=</span><span class="n">marker_size</span><span class="p">,</span>
                    <span class="n">facecolor</span><span class="o">=</span><span class="s2">&quot;w&quot;</span><span class="p">,</span>
                    <span class="n">edgecolor</span><span class="o">=</span><span class="n">facecolor</span><span class="p">,</span>
                    <span class="n">zorder</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span>
                    <span class="n">transform</span><span class="o">=</span><span class="n">transform</span><span class="p">,</span>
                <span class="p">)</span>
                <span class="n">ax</span><span class="o">.</span><span class="n">add_collection</span><span class="p">(</span><span class="n">circle_marker_start</span><span class="p">)</span>
            <span class="k">elif</span> <span class="n">d</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;outer_edge_type&quot;</span><span class="p">)</span> <span class="ow">in</span> <span class="p">[</span><span class="s2">&quot;+--&quot;</span><span class="p">,</span> <span class="s2">&quot;+-&gt;&quot;</span><span class="p">]:</span>
                <span class="n">circle_marker_start</span> <span class="o">=</span> <span class="n">ax</span><span class="o">.</span><span class="n">scatter</span><span class="p">(</span>
                    <span class="o">*</span><span class="n">start</span><span class="p">,</span>
                    <span class="n">marker</span><span class="o">=</span><span class="s2">&quot;P&quot;</span><span class="p">,</span>
                    <span class="n">s</span><span class="o">=</span><span class="n">marker_size</span><span class="p">,</span>
                    <span class="n">facecolor</span><span class="o">=</span><span class="s2">&quot;w&quot;</span><span class="p">,</span>
                    <span class="n">edgecolor</span><span class="o">=</span><span class="n">facecolor</span><span class="p">,</span>
                    <span class="n">zorder</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span>
                    <span class="n">transform</span><span class="o">=</span><span class="n">transform</span><span class="p">,</span>
                <span class="p">)</span>
                <span class="n">ax</span><span class="o">.</span><span class="n">add_collection</span><span class="p">(</span><span class="n">circle_marker_start</span><span class="p">)</span>
            <span class="k">elif</span> <span class="n">d</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;outer_edge_type&quot;</span><span class="p">)</span> <span class="o">==</span> <span class="s2">&quot;&lt;-x&quot;</span><span class="p">:</span>
                <span class="n">circle_marker_end</span> <span class="o">=</span> <span class="n">ax</span><span class="o">.</span><span class="n">scatter</span><span class="p">(</span>
                    <span class="o">*</span><span class="n">start</span><span class="p">,</span>
                    <span class="n">marker</span><span class="o">=</span><span class="s2">&quot;X&quot;</span><span class="p">,</span>
                    <span class="n">s</span><span class="o">=</span><span class="n">marker_size</span><span class="p">,</span>
                    <span class="n">facecolor</span><span class="o">=</span><span class="s2">&quot;w&quot;</span><span class="p">,</span>
                    <span class="n">edgecolor</span><span class="o">=</span><span class="n">facecolor</span><span class="p">,</span>
                    <span class="n">zorder</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span>
                    <span class="n">transform</span><span class="o">=</span><span class="n">transform</span><span class="p">,</span>
                <span class="p">)</span>
                <span class="n">ax</span><span class="o">.</span><span class="n">add_collection</span><span class="p">(</span><span class="n">circle_marker_end</span><span class="p">)</span>
            <span class="k">elif</span> <span class="n">d</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;outer_edge_type&quot;</span><span class="p">)</span> <span class="o">==</span> <span class="s2">&quot;&lt;-+&quot;</span><span class="p">:</span>
                <span class="n">circle_marker_end</span> <span class="o">=</span> <span class="n">ax</span><span class="o">.</span><span class="n">scatter</span><span class="p">(</span>
                    <span class="o">*</span><span class="n">start</span><span class="p">,</span>
                    <span class="n">marker</span><span class="o">=</span><span class="s2">&quot;P&quot;</span><span class="p">,</span>
                    <span class="n">s</span><span class="o">=</span><span class="n">marker_size</span><span class="p">,</span>
                    <span class="n">facecolor</span><span class="o">=</span><span class="s2">&quot;w&quot;</span><span class="p">,</span>
                    <span class="n">edgecolor</span><span class="o">=</span><span class="n">facecolor</span><span class="p">,</span>
                    <span class="n">zorder</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span>
                    <span class="n">transform</span><span class="o">=</span><span class="n">transform</span><span class="p">,</span>
                <span class="p">)</span>
                <span class="n">ax</span><span class="o">.</span><span class="n">add_collection</span><span class="p">(</span><span class="n">circle_marker_end</span><span class="p">)</span>
            <span class="k">elif</span> <span class="n">d</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;outer_edge_type&quot;</span><span class="p">)</span> <span class="o">==</span> <span class="s2">&quot;--x&quot;</span><span class="p">:</span>
                <span class="n">circle_marker_end</span> <span class="o">=</span> <span class="n">ax</span><span class="o">.</span><span class="n">scatter</span><span class="p">(</span>
                    <span class="o">*</span><span class="n">end</span><span class="p">,</span>
                    <span class="n">marker</span><span class="o">=</span><span class="s2">&quot;X&quot;</span><span class="p">,</span>
                    <span class="n">s</span><span class="o">=</span><span class="n">marker_size</span><span class="p">,</span>
                    <span class="n">facecolor</span><span class="o">=</span><span class="s2">&quot;w&quot;</span><span class="p">,</span>
                    <span class="n">edgecolor</span><span class="o">=</span><span class="n">facecolor</span><span class="p">,</span>
                    <span class="n">zorder</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span>
                    <span class="n">transform</span><span class="o">=</span><span class="n">transform</span><span class="p">,</span>
                <span class="p">)</span>
                <span class="n">ax</span><span class="o">.</span><span class="n">add_collection</span><span class="p">(</span><span class="n">circle_marker_end</span><span class="p">)</span>
            <span class="k">elif</span> <span class="n">d</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;outer_edge_type&quot;</span><span class="p">)</span> <span class="o">==</span> <span class="s2">&quot;o-o&quot;</span><span class="p">:</span>
                <span class="n">circle_marker_start</span> <span class="o">=</span> <span class="n">ax</span><span class="o">.</span><span class="n">scatter</span><span class="p">(</span>
                    <span class="o">*</span><span class="n">start</span><span class="p">,</span>
                    <span class="n">marker</span><span class="o">=</span><span class="s2">&quot;o&quot;</span><span class="p">,</span>
                    <span class="n">s</span><span class="o">=</span><span class="n">marker_size</span><span class="p">,</span>
                    <span class="n">facecolor</span><span class="o">=</span><span class="s2">&quot;w&quot;</span><span class="p">,</span>
                    <span class="n">edgecolor</span><span class="o">=</span><span class="n">facecolor</span><span class="p">,</span>
                    <span class="n">zorder</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span>
                    <span class="n">transform</span><span class="o">=</span><span class="n">transform</span><span class="p">,</span>
                <span class="p">)</span>
                <span class="n">ax</span><span class="o">.</span><span class="n">add_collection</span><span class="p">(</span><span class="n">circle_marker_start</span><span class="p">)</span>
                <span class="n">circle_marker_end</span> <span class="o">=</span> <span class="n">ax</span><span class="o">.</span><span class="n">scatter</span><span class="p">(</span>
                    <span class="o">*</span><span class="n">end</span><span class="p">,</span>
                    <span class="n">marker</span><span class="o">=</span><span class="s2">&quot;o&quot;</span><span class="p">,</span>
                    <span class="n">s</span><span class="o">=</span><span class="n">marker_size</span><span class="p">,</span>
                    <span class="n">facecolor</span><span class="o">=</span><span class="s2">&quot;w&quot;</span><span class="p">,</span>
                    <span class="n">edgecolor</span><span class="o">=</span><span class="n">facecolor</span><span class="p">,</span>
                    <span class="n">zorder</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span>
                    <span class="n">transform</span><span class="o">=</span><span class="n">transform</span><span class="p">,</span>
                <span class="p">)</span>
                <span class="n">ax</span><span class="o">.</span><span class="n">add_collection</span><span class="p">(</span><span class="n">circle_marker_end</span><span class="p">)</span>
            <span class="k">elif</span> <span class="n">d</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;outer_edge_type&quot;</span><span class="p">)</span> <span class="o">==</span> <span class="s2">&quot;x-x&quot;</span><span class="p">:</span>
                <span class="n">circle_marker_start</span> <span class="o">=</span> <span class="n">ax</span><span class="o">.</span><span class="n">scatter</span><span class="p">(</span>
                    <span class="o">*</span><span class="n">start</span><span class="p">,</span>
                    <span class="n">marker</span><span class="o">=</span><span class="s2">&quot;X&quot;</span><span class="p">,</span>
                    <span class="n">s</span><span class="o">=</span><span class="n">marker_size</span><span class="p">,</span>
                    <span class="n">facecolor</span><span class="o">=</span><span class="s2">&quot;w&quot;</span><span class="p">,</span>
                    <span class="n">edgecolor</span><span class="o">=</span><span class="n">facecolor</span><span class="p">,</span>
                    <span class="n">zorder</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span>
                    <span class="n">transform</span><span class="o">=</span><span class="n">transform</span><span class="p">,</span>
                <span class="p">)</span>
                <span class="n">ax</span><span class="o">.</span><span class="n">add_collection</span><span class="p">(</span><span class="n">circle_marker_start</span><span class="p">)</span>
                <span class="n">circle_marker_end</span> <span class="o">=</span> <span class="n">ax</span><span class="o">.</span><span class="n">scatter</span><span class="p">(</span>
                    <span class="o">*</span><span class="n">end</span><span class="p">,</span>
                    <span class="n">marker</span><span class="o">=</span><span class="s2">&quot;X&quot;</span><span class="p">,</span>
                    <span class="n">s</span><span class="o">=</span><span class="n">marker_size</span><span class="p">,</span>
                    <span class="n">facecolor</span><span class="o">=</span><span class="s2">&quot;w&quot;</span><span class="p">,</span>
                    <span class="n">edgecolor</span><span class="o">=</span><span class="n">facecolor</span><span class="p">,</span>
                    <span class="n">zorder</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span>
                    <span class="n">transform</span><span class="o">=</span><span class="n">transform</span><span class="p">,</span>
                <span class="p">)</span>
                <span class="n">ax</span><span class="o">.</span><span class="n">add_collection</span><span class="p">(</span><span class="n">circle_marker_end</span><span class="p">)</span>
            <span class="k">elif</span> <span class="n">d</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;outer_edge_type&quot;</span><span class="p">)</span> <span class="o">==</span> <span class="s2">&quot;o-x&quot;</span><span class="p">:</span>
                <span class="n">circle_marker_start</span> <span class="o">=</span> <span class="n">ax</span><span class="o">.</span><span class="n">scatter</span><span class="p">(</span>
                    <span class="o">*</span><span class="n">start</span><span class="p">,</span>
                    <span class="n">marker</span><span class="o">=</span><span class="s2">&quot;o&quot;</span><span class="p">,</span>
                    <span class="n">s</span><span class="o">=</span><span class="n">marker_size</span><span class="p">,</span>
                    <span class="n">facecolor</span><span class="o">=</span><span class="s2">&quot;w&quot;</span><span class="p">,</span>
                    <span class="n">edgecolor</span><span class="o">=</span><span class="n">facecolor</span><span class="p">,</span>
                    <span class="n">zorder</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span>
                    <span class="n">transform</span><span class="o">=</span><span class="n">transform</span><span class="p">,</span>
                <span class="p">)</span>
                <span class="n">ax</span><span class="o">.</span><span class="n">add_collection</span><span class="p">(</span><span class="n">circle_marker_start</span><span class="p">)</span>
                <span class="n">circle_marker_end</span> <span class="o">=</span> <span class="n">ax</span><span class="o">.</span><span class="n">scatter</span><span class="p">(</span>
                    <span class="o">*</span><span class="n">end</span><span class="p">,</span>
                    <span class="n">marker</span><span class="o">=</span><span class="s2">&quot;X&quot;</span><span class="p">,</span>
                    <span class="n">s</span><span class="o">=</span><span class="n">marker_size</span><span class="p">,</span>
                    <span class="n">facecolor</span><span class="o">=</span><span class="s2">&quot;w&quot;</span><span class="p">,</span>
                    <span class="n">edgecolor</span><span class="o">=</span><span class="n">facecolor</span><span class="p">,</span>
                    <span class="n">zorder</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span>
                    <span class="n">transform</span><span class="o">=</span><span class="n">transform</span><span class="p">,</span>
                <span class="p">)</span>
                <span class="n">ax</span><span class="o">.</span><span class="n">add_collection</span><span class="p">(</span><span class="n">circle_marker_end</span><span class="p">)</span>
            <span class="k">elif</span> <span class="n">d</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;outer_edge_type&quot;</span><span class="p">)</span> <span class="o">==</span> <span class="s2">&quot;x-o&quot;</span><span class="p">:</span>
                <span class="n">circle_marker_start</span> <span class="o">=</span> <span class="n">ax</span><span class="o">.</span><span class="n">scatter</span><span class="p">(</span>
                    <span class="o">*</span><span class="n">start</span><span class="p">,</span>
                    <span class="n">marker</span><span class="o">=</span><span class="s2">&quot;X&quot;</span><span class="p">,</span>
                    <span class="n">s</span><span class="o">=</span><span class="n">marker_size</span><span class="p">,</span>
                    <span class="n">facecolor</span><span class="o">=</span><span class="s2">&quot;w&quot;</span><span class="p">,</span>
                    <span class="n">edgecolor</span><span class="o">=</span><span class="n">facecolor</span><span class="p">,</span>
                    <span class="n">zorder</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span>
                    <span class="n">transform</span><span class="o">=</span><span class="n">transform</span><span class="p">,</span>
                <span class="p">)</span>
                <span class="n">ax</span><span class="o">.</span><span class="n">add_collection</span><span class="p">(</span><span class="n">circle_marker_start</span><span class="p">)</span>
                <span class="n">circle_marker_end</span> <span class="o">=</span> <span class="n">ax</span><span class="o">.</span><span class="n">scatter</span><span class="p">(</span>
                    <span class="o">*</span><span class="n">end</span><span class="p">,</span>
                    <span class="n">marker</span><span class="o">=</span><span class="s2">&quot;o&quot;</span><span class="p">,</span>
                    <span class="n">s</span><span class="o">=</span><span class="n">marker_size</span><span class="p">,</span>
                    <span class="n">facecolor</span><span class="o">=</span><span class="s2">&quot;w&quot;</span><span class="p">,</span>
                    <span class="n">edgecolor</span><span class="o">=</span><span class="n">facecolor</span><span class="p">,</span>
                    <span class="n">zorder</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span>
                    <span class="n">transform</span><span class="o">=</span><span class="n">transform</span><span class="p">,</span>
                <span class="p">)</span>
                <span class="n">ax</span><span class="o">.</span><span class="n">add_collection</span><span class="p">(</span><span class="n">circle_marker_end</span><span class="p">)</span>

        <span class="k">else</span><span class="p">:</span>
            <span class="k">if</span> <span class="n">d</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;inner_edge_type&quot;</span><span class="p">)</span> <span class="ow">in</span> <span class="p">[</span><span class="s2">&quot;o-&gt;&quot;</span><span class="p">,</span> <span class="s2">&quot;o--&quot;</span><span class="p">]:</span>
                <span class="n">circle_marker_start</span> <span class="o">=</span> <span class="n">ax</span><span class="o">.</span><span class="n">scatter</span><span class="p">(</span>
                    <span class="o">*</span><span class="n">start</span><span class="p">,</span>
                    <span class="n">marker</span><span class="o">=</span><span class="s2">&quot;o&quot;</span><span class="p">,</span>
                    <span class="n">s</span><span class="o">=</span><span class="n">marker_size</span><span class="p">,</span>
                    <span class="n">facecolor</span><span class="o">=</span><span class="s2">&quot;w&quot;</span><span class="p">,</span>
                    <span class="n">edgecolor</span><span class="o">=</span><span class="n">facecolor</span><span class="p">,</span>
                    <span class="n">zorder</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span>
                    <span class="n">transform</span><span class="o">=</span><span class="n">transform</span><span class="p">,</span>
                <span class="p">)</span>
                <span class="n">ax</span><span class="o">.</span><span class="n">add_collection</span><span class="p">(</span><span class="n">circle_marker_start</span><span class="p">)</span>
            <span class="k">elif</span> <span class="n">d</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;inner_edge_type&quot;</span><span class="p">)</span> <span class="o">==</span> <span class="s2">&quot;&lt;-o&quot;</span><span class="p">:</span>
                <span class="n">circle_marker_end</span> <span class="o">=</span> <span class="n">ax</span><span class="o">.</span><span class="n">scatter</span><span class="p">(</span>
                    <span class="o">*</span><span class="n">start</span><span class="p">,</span>
                    <span class="n">marker</span><span class="o">=</span><span class="s2">&quot;o&quot;</span><span class="p">,</span>
                    <span class="n">s</span><span class="o">=</span><span class="n">marker_size</span><span class="p">,</span>
                    <span class="n">facecolor</span><span class="o">=</span><span class="s2">&quot;w&quot;</span><span class="p">,</span>
                    <span class="n">edgecolor</span><span class="o">=</span><span class="n">facecolor</span><span class="p">,</span>
                    <span class="n">zorder</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span>
                    <span class="n">transform</span><span class="o">=</span><span class="n">transform</span><span class="p">,</span>
                <span class="p">)</span>
                <span class="n">ax</span><span class="o">.</span><span class="n">add_collection</span><span class="p">(</span><span class="n">circle_marker_end</span><span class="p">)</span>
            <span class="k">elif</span> <span class="n">d</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;inner_edge_type&quot;</span><span class="p">)</span> <span class="o">==</span> <span class="s2">&quot;--o&quot;</span><span class="p">:</span>
                <span class="n">circle_marker_end</span> <span class="o">=</span> <span class="n">ax</span><span class="o">.</span><span class="n">scatter</span><span class="p">(</span>
                    <span class="o">*</span><span class="n">end</span><span class="p">,</span>
                    <span class="n">marker</span><span class="o">=</span><span class="s2">&quot;o&quot;</span><span class="p">,</span>
                    <span class="n">s</span><span class="o">=</span><span class="n">marker_size</span><span class="p">,</span>
                    <span class="n">facecolor</span><span class="o">=</span><span class="s2">&quot;w&quot;</span><span class="p">,</span>
                    <span class="n">edgecolor</span><span class="o">=</span><span class="n">facecolor</span><span class="p">,</span>
                    <span class="n">zorder</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span>
                    <span class="n">transform</span><span class="o">=</span><span class="n">transform</span><span class="p">,</span>
                <span class="p">)</span>
                <span class="n">ax</span><span class="o">.</span><span class="n">add_collection</span><span class="p">(</span><span class="n">circle_marker_end</span><span class="p">)</span>
            <span class="k">elif</span> <span class="n">d</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;inner_edge_type&quot;</span><span class="p">)</span> <span class="ow">in</span> <span class="p">[</span><span class="s2">&quot;x--&quot;</span><span class="p">,</span> <span class="s2">&quot;x-&gt;&quot;</span><span class="p">]:</span>
                <span class="n">circle_marker_start</span> <span class="o">=</span> <span class="n">ax</span><span class="o">.</span><span class="n">scatter</span><span class="p">(</span>
                    <span class="o">*</span><span class="n">start</span><span class="p">,</span>
                    <span class="n">marker</span><span class="o">=</span><span class="s2">&quot;X&quot;</span><span class="p">,</span>
                    <span class="n">s</span><span class="o">=</span><span class="n">marker_size</span><span class="p">,</span>
                    <span class="n">facecolor</span><span class="o">=</span><span class="s2">&quot;w&quot;</span><span class="p">,</span>
                    <span class="n">edgecolor</span><span class="o">=</span><span class="n">facecolor</span><span class="p">,</span>
                    <span class="n">zorder</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span>
                    <span class="n">transform</span><span class="o">=</span><span class="n">transform</span><span class="p">,</span>
                <span class="p">)</span>
                <span class="n">ax</span><span class="o">.</span><span class="n">add_collection</span><span class="p">(</span><span class="n">circle_marker_start</span><span class="p">)</span>
            <span class="k">elif</span> <span class="n">d</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;inner_edge_type&quot;</span><span class="p">)</span> <span class="ow">in</span> <span class="p">[</span><span class="s2">&quot;+--&quot;</span><span class="p">,</span> <span class="s2">&quot;+-&gt;&quot;</span><span class="p">]:</span>
                <span class="n">circle_marker_start</span> <span class="o">=</span> <span class="n">ax</span><span class="o">.</span><span class="n">scatter</span><span class="p">(</span>
                    <span class="o">*</span><span class="n">start</span><span class="p">,</span>
                    <span class="n">marker</span><span class="o">=</span><span class="s2">&quot;P&quot;</span><span class="p">,</span>
                    <span class="n">s</span><span class="o">=</span><span class="n">marker_size</span><span class="p">,</span>
                    <span class="n">facecolor</span><span class="o">=</span><span class="s2">&quot;w&quot;</span><span class="p">,</span>
                    <span class="n">edgecolor</span><span class="o">=</span><span class="n">facecolor</span><span class="p">,</span>
                    <span class="n">zorder</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span>
                    <span class="n">transform</span><span class="o">=</span><span class="n">transform</span><span class="p">,</span>
                <span class="p">)</span>
                <span class="n">ax</span><span class="o">.</span><span class="n">add_collection</span><span class="p">(</span><span class="n">circle_marker_start</span><span class="p">)</span>
            <span class="k">elif</span> <span class="n">d</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;inner_edge_type&quot;</span><span class="p">)</span> <span class="o">==</span> <span class="s2">&quot;&lt;-x&quot;</span><span class="p">:</span>
                <span class="n">circle_marker_end</span> <span class="o">=</span> <span class="n">ax</span><span class="o">.</span><span class="n">scatter</span><span class="p">(</span>
                    <span class="o">*</span><span class="n">start</span><span class="p">,</span>
                    <span class="n">marker</span><span class="o">=</span><span class="s2">&quot;X&quot;</span><span class="p">,</span>
                    <span class="n">s</span><span class="o">=</span><span class="n">marker_size</span><span class="p">,</span>
                    <span class="n">facecolor</span><span class="o">=</span><span class="s2">&quot;w&quot;</span><span class="p">,</span>
                    <span class="n">edgecolor</span><span class="o">=</span><span class="n">facecolor</span><span class="p">,</span>
                    <span class="n">zorder</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span>
                    <span class="n">transform</span><span class="o">=</span><span class="n">transform</span><span class="p">,</span>
                <span class="p">)</span>
                <span class="n">ax</span><span class="o">.</span><span class="n">add_collection</span><span class="p">(</span><span class="n">circle_marker_end</span><span class="p">)</span>
            <span class="k">elif</span> <span class="n">d</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;inner_edge_type&quot;</span><span class="p">)</span> <span class="o">==</span> <span class="s2">&quot;&lt;-+&quot;</span><span class="p">:</span>
                <span class="n">circle_marker_end</span> <span class="o">=</span> <span class="n">ax</span><span class="o">.</span><span class="n">scatter</span><span class="p">(</span>
                    <span class="o">*</span><span class="n">start</span><span class="p">,</span>
                    <span class="n">marker</span><span class="o">=</span><span class="s2">&quot;P&quot;</span><span class="p">,</span>
                    <span class="n">s</span><span class="o">=</span><span class="n">marker_size</span><span class="p">,</span>
                    <span class="n">facecolor</span><span class="o">=</span><span class="s2">&quot;w&quot;</span><span class="p">,</span>
                    <span class="n">edgecolor</span><span class="o">=</span><span class="n">facecolor</span><span class="p">,</span>
                    <span class="n">zorder</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span>
                    <span class="n">transform</span><span class="o">=</span><span class="n">transform</span><span class="p">,</span>
                <span class="p">)</span>
                <span class="n">ax</span><span class="o">.</span><span class="n">add_collection</span><span class="p">(</span><span class="n">circle_marker_end</span><span class="p">)</span>
            <span class="k">elif</span> <span class="n">d</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;inner_edge_type&quot;</span><span class="p">)</span> <span class="o">==</span> <span class="s2">&quot;--x&quot;</span><span class="p">:</span>
                <span class="n">circle_marker_end</span> <span class="o">=</span> <span class="n">ax</span><span class="o">.</span><span class="n">scatter</span><span class="p">(</span>
                    <span class="o">*</span><span class="n">end</span><span class="p">,</span>
                    <span class="n">marker</span><span class="o">=</span><span class="s2">&quot;X&quot;</span><span class="p">,</span>
                    <span class="n">s</span><span class="o">=</span><span class="n">marker_size</span><span class="p">,</span>
                    <span class="n">facecolor</span><span class="o">=</span><span class="s2">&quot;w&quot;</span><span class="p">,</span>
                    <span class="n">edgecolor</span><span class="o">=</span><span class="n">facecolor</span><span class="p">,</span>
                    <span class="n">zorder</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span>
                    <span class="n">transform</span><span class="o">=</span><span class="n">transform</span><span class="p">,</span>
                <span class="p">)</span>
                <span class="n">ax</span><span class="o">.</span><span class="n">add_collection</span><span class="p">(</span><span class="n">circle_marker_end</span><span class="p">)</span>
            <span class="k">elif</span> <span class="n">d</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;inner_edge_type&quot;</span><span class="p">)</span> <span class="o">==</span> <span class="s2">&quot;o-o&quot;</span><span class="p">:</span>
                <span class="n">circle_marker_start</span> <span class="o">=</span> <span class="n">ax</span><span class="o">.</span><span class="n">scatter</span><span class="p">(</span>
                    <span class="o">*</span><span class="n">start</span><span class="p">,</span>
                    <span class="n">marker</span><span class="o">=</span><span class="s2">&quot;o&quot;</span><span class="p">,</span>
                    <span class="n">s</span><span class="o">=</span><span class="n">marker_size</span><span class="p">,</span>
                    <span class="n">facecolor</span><span class="o">=</span><span class="s2">&quot;w&quot;</span><span class="p">,</span>
                    <span class="n">edgecolor</span><span class="o">=</span><span class="n">facecolor</span><span class="p">,</span>
                    <span class="n">zorder</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span>
                    <span class="n">transform</span><span class="o">=</span><span class="n">transform</span><span class="p">,</span>
                <span class="p">)</span>
                <span class="n">ax</span><span class="o">.</span><span class="n">add_collection</span><span class="p">(</span><span class="n">circle_marker_start</span><span class="p">)</span>
                <span class="n">circle_marker_end</span> <span class="o">=</span> <span class="n">ax</span><span class="o">.</span><span class="n">scatter</span><span class="p">(</span>
                    <span class="o">*</span><span class="n">end</span><span class="p">,</span>
                    <span class="n">marker</span><span class="o">=</span><span class="s2">&quot;o&quot;</span><span class="p">,</span>
                    <span class="n">s</span><span class="o">=</span><span class="n">marker_size</span><span class="p">,</span>
                    <span class="n">facecolor</span><span class="o">=</span><span class="s2">&quot;w&quot;</span><span class="p">,</span>
                    <span class="n">edgecolor</span><span class="o">=</span><span class="n">facecolor</span><span class="p">,</span>
                    <span class="n">zorder</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span>
                    <span class="n">transform</span><span class="o">=</span><span class="n">transform</span><span class="p">,</span>
                <span class="p">)</span>
                <span class="n">ax</span><span class="o">.</span><span class="n">add_collection</span><span class="p">(</span><span class="n">circle_marker_end</span><span class="p">)</span>
            <span class="k">elif</span> <span class="n">d</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;inner_edge_type&quot;</span><span class="p">)</span> <span class="o">==</span> <span class="s2">&quot;x-x&quot;</span><span class="p">:</span>
                <span class="n">circle_marker_start</span> <span class="o">=</span> <span class="n">ax</span><span class="o">.</span><span class="n">scatter</span><span class="p">(</span>
                    <span class="o">*</span><span class="n">start</span><span class="p">,</span>
                    <span class="n">marker</span><span class="o">=</span><span class="s2">&quot;X&quot;</span><span class="p">,</span>
                    <span class="n">s</span><span class="o">=</span><span class="n">marker_size</span><span class="p">,</span>
                    <span class="n">facecolor</span><span class="o">=</span><span class="s2">&quot;w&quot;</span><span class="p">,</span>
                    <span class="n">edgecolor</span><span class="o">=</span><span class="n">facecolor</span><span class="p">,</span>
                    <span class="n">zorder</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span>
                    <span class="n">transform</span><span class="o">=</span><span class="n">transform</span><span class="p">,</span>
                <span class="p">)</span>
                <span class="n">ax</span><span class="o">.</span><span class="n">add_collection</span><span class="p">(</span><span class="n">circle_marker_start</span><span class="p">)</span>
                <span class="n">circle_marker_end</span> <span class="o">=</span> <span class="n">ax</span><span class="o">.</span><span class="n">scatter</span><span class="p">(</span>
                    <span class="o">*</span><span class="n">end</span><span class="p">,</span>
                    <span class="n">marker</span><span class="o">=</span><span class="s2">&quot;X&quot;</span><span class="p">,</span>
                    <span class="n">s</span><span class="o">=</span><span class="n">marker_size</span><span class="p">,</span>
                    <span class="n">facecolor</span><span class="o">=</span><span class="s2">&quot;w&quot;</span><span class="p">,</span>
                    <span class="n">edgecolor</span><span class="o">=</span><span class="n">facecolor</span><span class="p">,</span>
                    <span class="n">zorder</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span>
                    <span class="n">transform</span><span class="o">=</span><span class="n">transform</span><span class="p">,</span>
                <span class="p">)</span>
                <span class="n">ax</span><span class="o">.</span><span class="n">add_collection</span><span class="p">(</span><span class="n">circle_marker_end</span><span class="p">)</span>
            <span class="k">elif</span> <span class="n">d</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;inner_edge_type&quot;</span><span class="p">)</span> <span class="o">==</span> <span class="s2">&quot;o-x&quot;</span><span class="p">:</span>
                <span class="n">circle_marker_start</span> <span class="o">=</span> <span class="n">ax</span><span class="o">.</span><span class="n">scatter</span><span class="p">(</span>
                    <span class="o">*</span><span class="n">start</span><span class="p">,</span>
                    <span class="n">marker</span><span class="o">=</span><span class="s2">&quot;o&quot;</span><span class="p">,</span>
                    <span class="n">s</span><span class="o">=</span><span class="n">marker_size</span><span class="p">,</span>
                    <span class="n">facecolor</span><span class="o">=</span><span class="s2">&quot;w&quot;</span><span class="p">,</span>
                    <span class="n">edgecolor</span><span class="o">=</span><span class="n">facecolor</span><span class="p">,</span>
                    <span class="n">zorder</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span>
                    <span class="n">transform</span><span class="o">=</span><span class="n">transform</span><span class="p">,</span>
                <span class="p">)</span>
                <span class="n">ax</span><span class="o">.</span><span class="n">add_collection</span><span class="p">(</span><span class="n">circle_marker_start</span><span class="p">)</span>
                <span class="n">circle_marker_end</span> <span class="o">=</span> <span class="n">ax</span><span class="o">.</span><span class="n">scatter</span><span class="p">(</span>
                    <span class="o">*</span><span class="n">end</span><span class="p">,</span>
                    <span class="n">marker</span><span class="o">=</span><span class="s2">&quot;X&quot;</span><span class="p">,</span>
                    <span class="n">s</span><span class="o">=</span><span class="n">marker_size</span><span class="p">,</span>
                    <span class="n">facecolor</span><span class="o">=</span><span class="s2">&quot;w&quot;</span><span class="p">,</span>
                    <span class="n">edgecolor</span><span class="o">=</span><span class="n">facecolor</span><span class="p">,</span>
                    <span class="n">zorder</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span>
                    <span class="n">transform</span><span class="o">=</span><span class="n">transform</span><span class="p">,</span>
                <span class="p">)</span>
                <span class="n">ax</span><span class="o">.</span><span class="n">add_collection</span><span class="p">(</span><span class="n">circle_marker_end</span><span class="p">)</span>
            <span class="k">elif</span> <span class="n">d</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;inner_edge_type&quot;</span><span class="p">)</span> <span class="o">==</span> <span class="s2">&quot;x-o&quot;</span><span class="p">:</span>
                <span class="n">circle_marker_start</span> <span class="o">=</span> <span class="n">ax</span><span class="o">.</span><span class="n">scatter</span><span class="p">(</span>
                    <span class="o">*</span><span class="n">start</span><span class="p">,</span>
                    <span class="n">marker</span><span class="o">=</span><span class="s2">&quot;X&quot;</span><span class="p">,</span>
                    <span class="n">s</span><span class="o">=</span><span class="n">marker_size</span><span class="p">,</span>
                    <span class="n">facecolor</span><span class="o">=</span><span class="s2">&quot;w&quot;</span><span class="p">,</span>
                    <span class="n">edgecolor</span><span class="o">=</span><span class="n">facecolor</span><span class="p">,</span>
                    <span class="n">zorder</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span>
                    <span class="n">transform</span><span class="o">=</span><span class="n">transform</span><span class="p">,</span>
                <span class="p">)</span>
                <span class="n">ax</span><span class="o">.</span><span class="n">add_collection</span><span class="p">(</span><span class="n">circle_marker_start</span><span class="p">)</span>
                <span class="n">circle_marker_end</span> <span class="o">=</span> <span class="n">ax</span><span class="o">.</span><span class="n">scatter</span><span class="p">(</span>
                    <span class="o">*</span><span class="n">end</span><span class="p">,</span>
                    <span class="n">marker</span><span class="o">=</span><span class="s2">&quot;o&quot;</span><span class="p">,</span>
                    <span class="n">s</span><span class="o">=</span><span class="n">marker_size</span><span class="p">,</span>
                    <span class="n">facecolor</span><span class="o">=</span><span class="s2">&quot;w&quot;</span><span class="p">,</span>
                    <span class="n">edgecolor</span><span class="o">=</span><span class="n">facecolor</span><span class="p">,</span>
                    <span class="n">zorder</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span>
                    <span class="n">transform</span><span class="o">=</span><span class="n">transform</span><span class="p">,</span>
                <span class="p">)</span>
                <span class="n">ax</span><span class="o">.</span><span class="n">add_collection</span><span class="p">(</span><span class="n">circle_marker_end</span><span class="p">)</span>



        <span class="k">if</span> <span class="n">d</span><span class="p">[</span><span class="s2">&quot;label&quot;</span><span class="p">]</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> <span class="n">outer_edge</span><span class="p">:</span>
            <span class="k">def</span> <span class="nf">closest_node</span><span class="p">(</span><span class="n">node</span><span class="p">,</span> <span class="n">nodes</span><span class="p">):</span>
                <span class="n">nodes</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">asarray</span><span class="p">(</span><span class="n">nodes</span><span class="p">)</span>
                <span class="n">node</span> <span class="o">=</span> <span class="n">node</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">)</span>
                <span class="n">dist_2</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">sum</span><span class="p">((</span><span class="n">nodes</span> <span class="o">-</span> <span class="n">node</span><span class="p">)</span><span class="o">**</span><span class="mi">2</span><span class="p">,</span> <span class="n">axis</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span>
                <span class="k">return</span> <span class="n">np</span><span class="o">.</span><span class="n">argmin</span><span class="p">(</span><span class="n">dist_2</span><span class="p">)</span>

            <span class="c1"># Attach labels of lags</span>
            <span class="c1"># trans = None  # patch.get_transform()</span>
            <span class="c1"># path = e_p.get_path()</span>
            <span class="n">vertices</span> <span class="o">=</span> <span class="n">e_p_marker</span><span class="o">.</span><span class="n">get_path</span><span class="p">()</span><span class="o">.</span><span class="n">vertices</span><span class="o">.</span><span class="n">copy</span><span class="p">()</span>
            <span class="n">verts</span> <span class="o">=</span> <span class="n">e_p</span><span class="o">.</span><span class="n">get_path</span><span class="p">()</span><span class="o">.</span><span class="n">to_polygons</span><span class="p">(</span><span class="n">transform</span><span class="o">=</span><span class="kc">None</span><span class="p">)[</span><span class="mi">0</span><span class="p">]</span>
            <span class="c1"># print(verts)</span>
            <span class="c1"># print(verts.shape)</span>
            <span class="c1"># print(vertices.shape)</span>
            <span class="c1"># for num, vert in enumerate(verts):</span>
            <span class="c1">#     ax.text(vert[0], vert[1], str(num), </span>
            <span class="c1">#         transform=transform,)</span>
            <span class="c1"># ax.scatter(verts[:,0], verts[:,1])</span>
            <span class="c1"># mid_point = np.array([(start[0] + end[0])/2., (start[1] + end[1])/2.])</span>
            <span class="c1"># print(start, end, mid_point)</span>
            <span class="c1"># ax.scatter(mid_point[0], mid_point[1], marker=&#39;x&#39;, </span>
            <span class="c1">#     s=100, zorder=10, transform=transform,)</span>
            <span class="n">closest_node</span> <span class="o">=</span> <span class="n">closest_node</span><span class="p">(</span><span class="n">vertices</span><span class="p">[</span><span class="nb">int</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">vertices</span><span class="p">)</span><span class="o">/</span><span class="mf">2.</span><span class="p">),:],</span> <span class="n">verts</span><span class="p">)</span>
            <span class="c1"># print(closest_node, verts[closest_node])</span>
            <span class="c1"># ax.scatter(verts[closest_node][0], verts[closest_node][1], marker=&#39;x&#39;)</span>

            <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">vertices</span><span class="p">)</span> <span class="o">&gt;</span> <span class="mi">2</span><span class="p">:</span>
                <span class="c1"># label_vert = vertices[int(len(vertices)/2.),:] #verts[1, :]</span>
                <span class="n">label_vert</span> <span class="o">=</span> <span class="n">verts</span><span class="p">[</span><span class="n">closest_node</span><span class="p">]</span> <span class="c1">#verts[1, :]</span>
                <span class="n">l</span> <span class="o">=</span> <span class="n">d</span><span class="p">[</span><span class="s2">&quot;label&quot;</span><span class="p">]</span>
                <span class="n">string</span> <span class="o">=</span> <span class="nb">str</span><span class="p">(</span><span class="n">l</span><span class="p">)</span>
                <span class="n">txt</span> <span class="o">=</span> <span class="n">ax</span><span class="o">.</span><span class="n">text</span><span class="p">(</span>
                    <span class="n">label_vert</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span>
                    <span class="n">label_vert</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span>
                    <span class="n">string</span><span class="p">,</span>
                    <span class="n">fontsize</span><span class="o">=</span><span class="n">link_label_fontsize</span><span class="p">,</span>
                    <span class="n">verticalalignment</span><span class="o">=</span><span class="s2">&quot;center&quot;</span><span class="p">,</span>
                    <span class="n">horizontalalignment</span><span class="o">=</span><span class="s2">&quot;center&quot;</span><span class="p">,</span>
                    <span class="n">color</span><span class="o">=</span><span class="s2">&quot;w&quot;</span><span class="p">,</span>
                    <span class="n">zorder</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span>
                    <span class="n">transform</span><span class="o">=</span><span class="n">transform</span><span class="p">,</span>
                <span class="p">)</span>
                <span class="n">txt</span><span class="o">.</span><span class="n">set_path_effects</span><span class="p">(</span>
                    <span class="p">[</span><span class="n">PathEffects</span><span class="o">.</span><span class="n">withStroke</span><span class="p">(</span><span class="n">linewidth</span><span class="o">=</span><span class="mi">2</span><span class="p">,</span> <span class="n">foreground</span><span class="o">=</span><span class="s2">&quot;k&quot;</span><span class="p">)]</span>
                <span class="p">)</span>

        <span class="k">return</span> <span class="n">rad</span>

    <span class="c1"># Collect all edge weights to get color scale</span>
    <span class="n">all_links_weights</span> <span class="o">=</span> <span class="p">[]</span>
    <span class="n">all_links_edge_weights</span> <span class="o">=</span> <span class="p">[]</span>
    <span class="k">for</span> <span class="p">(</span><span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">,</span> <span class="n">d</span><span class="p">)</span> <span class="ow">in</span> <span class="n">G</span><span class="o">.</span><span class="n">edges</span><span class="p">(</span><span class="n">data</span><span class="o">=</span><span class="kc">True</span><span class="p">):</span>
        <span class="k">if</span> <span class="n">u</span> <span class="o">!=</span> <span class="n">v</span><span class="p">:</span>
            <span class="k">if</span> <span class="n">d</span><span class="p">[</span><span class="s2">&quot;outer_edge&quot;</span><span class="p">]</span> <span class="ow">and</span> <span class="n">d</span><span class="p">[</span><span class="s2">&quot;outer_edge_color&quot;</span><span class="p">]</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                <span class="n">all_links_weights</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">d</span><span class="p">[</span><span class="s2">&quot;outer_edge_color&quot;</span><span class="p">])</span>
            <span class="k">if</span> <span class="n">d</span><span class="p">[</span><span class="s2">&quot;inner_edge&quot;</span><span class="p">]</span> <span class="ow">and</span> <span class="n">d</span><span class="p">[</span><span class="s2">&quot;inner_edge_color&quot;</span><span class="p">]</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                <span class="n">all_links_weights</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">d</span><span class="p">[</span><span class="s2">&quot;inner_edge_color&quot;</span><span class="p">])</span>

    <span class="k">if</span> <span class="n">cmap_links</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> <span class="nb">len</span><span class="p">(</span><span class="n">all_links_weights</span><span class="p">)</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
        <span class="k">if</span> <span class="n">links_vmin</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">links_vmin</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">(</span><span class="n">all_links_weights</span><span class="p">)</span><span class="o">.</span><span class="n">min</span><span class="p">()</span>
        <span class="k">if</span> <span class="n">links_vmax</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">links_vmax</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">(</span><span class="n">all_links_weights</span><span class="p">)</span><span class="o">.</span><span class="n">max</span><span class="p">()</span>
        <span class="n">data_to_rgb_links</span> <span class="o">=</span> <span class="n">pyplot</span><span class="o">.</span><span class="n">cm</span><span class="o">.</span><span class="n">ScalarMappable</span><span class="p">(</span>
            <span class="n">norm</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">cmap</span><span class="o">=</span><span class="n">pyplot</span><span class="o">.</span><span class="n">get_cmap</span><span class="p">(</span><span class="n">cmap_links</span><span class="p">)</span>
        <span class="p">)</span>
        <span class="n">data_to_rgb_links</span><span class="o">.</span><span class="n">set_array</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">(</span><span class="n">all_links_weights</span><span class="p">))</span>
        <span class="n">data_to_rgb_links</span><span class="o">.</span><span class="n">set_clim</span><span class="p">(</span><span class="n">vmin</span><span class="o">=</span><span class="n">links_vmin</span><span class="p">,</span> <span class="n">vmax</span><span class="o">=</span><span class="n">links_vmax</span><span class="p">)</span>
        <span class="c1"># Create colorbars for links</span>

        <span class="c1"># setup colorbar axes.</span>
        <span class="k">if</span> <span class="n">show_colorbar</span><span class="p">:</span>
            <span class="c1"># cax_e = pyplot.axes(</span>
            <span class="c1">#     [</span>
            <span class="c1">#         0.55,</span>
            <span class="c1">#         ax.get_subplotspec().get_position(ax.figure).bounds[1] + 0.02,</span>
            <span class="c1">#         0.4,</span>
            <span class="c1">#         0.025 + (len(all_links_edge_weights) == 0) * 0.035,</span>
            <span class="c1">#     ],</span>
            <span class="c1">#     frameon=False,</span>
            <span class="c1"># )</span>
            <span class="n">bbox_ax</span> <span class="o">=</span> <span class="n">ax</span><span class="o">.</span><span class="n">get_position</span><span class="p">()</span>
            <span class="n">width</span> <span class="o">=</span> <span class="n">bbox_ax</span><span class="o">.</span><span class="n">xmax</span><span class="o">-</span><span class="n">bbox_ax</span><span class="o">.</span><span class="n">xmin</span>
            <span class="n">height</span> <span class="o">=</span> <span class="n">bbox_ax</span><span class="o">.</span><span class="n">ymax</span><span class="o">-</span><span class="n">bbox_ax</span><span class="o">.</span><span class="n">ymin</span>
            <span class="c1"># print(bbox_ax.xmin, bbox_ax.xmax, bbox_ax.ymin, bbox_ax.ymax) </span>
            <span class="c1"># cax_e = fig.add_axes(</span>
            <span class="c1">#     [</span>
            <span class="c1">#         bbox_ax.xmax - width*0.45,</span>
            <span class="c1">#         bbox_ax.ymin-0.075*height+network_lower_bound-0.15,</span>
            <span class="c1">#         width*0.4,</span>
            <span class="c1">#         0.075*height,   #0.025 + (len(all_links_edge_weights) == 0) * 0.035,</span>
            <span class="c1">#     ],</span>
            <span class="c1">#     frameon=False,</span>
            <span class="c1"># )</span>
            <span class="n">cax_e</span> <span class="o">=</span> <span class="n">ax</span><span class="o">.</span><span class="n">inset_axes</span><span class="p">(</span> 
                          <span class="p">[</span>
                          <span class="mf">0.55</span><span class="p">,</span> <span class="o">-</span><span class="mf">0.07</span><span class="p">,</span> <span class="mf">0.4</span><span class="p">,</span> <span class="mf">0.07</span>
                    <span class="c1"># bbox_ax.xmax - width*0.45,</span>
                    <span class="c1"># bbox_ax.ymin-0.075*height+network_lower_bound-0.15,</span>
                    <span class="c1"># width*0.4,</span>
                    <span class="c1"># 0.075*height,   #0.025 + (len(all_links_edge_weights) == 0) * 0.035,</span>
                <span class="p">],</span>
                <span class="n">frameon</span><span class="o">=</span><span class="kc">False</span><span class="p">,)</span>
            <span class="c1"># divider = make_axes_locatable(ax)</span>

            <span class="c1"># cax_e = divider.append_axes(&#39;bottom&#39;, size=&#39;5%&#39;, pad=0.05, frameon=False,)</span>

            <span class="n">cb_e</span> <span class="o">=</span> <span class="n">pyplot</span><span class="o">.</span><span class="n">colorbar</span><span class="p">(</span>
                <span class="n">data_to_rgb_links</span><span class="p">,</span> <span class="n">cax</span><span class="o">=</span><span class="n">cax_e</span><span class="p">,</span> <span class="n">orientation</span><span class="o">=</span><span class="s2">&quot;horizontal&quot;</span>
            <span class="p">)</span>
            <span class="c1"># try:</span>
            <span class="n">ticks_here</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span>
                    <span class="n">_myround</span><span class="p">(</span><span class="n">links_vmin</span><span class="p">,</span> <span class="n">links_ticks</span><span class="p">,</span> <span class="s2">&quot;down&quot;</span><span class="p">),</span>
                    <span class="n">_myround</span><span class="p">(</span><span class="n">links_vmax</span><span class="p">,</span> <span class="n">links_ticks</span><span class="p">,</span> <span class="s2">&quot;up&quot;</span><span class="p">)</span> <span class="o">+</span> <span class="n">links_ticks</span><span class="p">,</span>
                    <span class="n">links_ticks</span><span class="p">,</span>
                <span class="p">)</span>
            <span class="n">cb_e</span><span class="o">.</span><span class="n">set_ticks</span><span class="p">(</span><span class="n">ticks_here</span><span class="p">[(</span><span class="n">links_vmin</span> <span class="o">&lt;=</span> <span class="n">ticks_here</span><span class="p">)</span> <span class="o">&amp;</span> <span class="p">(</span><span class="n">ticks_here</span> <span class="o">&lt;=</span> <span class="n">links_vmax</span><span class="p">)])</span>
            <span class="c1"># except:</span>
            <span class="c1">#     print(&#39;no ticks given&#39;)</span>

            <span class="n">cb_e</span><span class="o">.</span><span class="n">outline</span><span class="o">.</span><span class="n">clear</span><span class="p">()</span>
            <span class="n">cax_e</span><span class="o">.</span><span class="n">set_xlabel</span><span class="p">(</span>
                <span class="n">link_colorbar_label</span><span class="p">,</span> <span class="n">labelpad</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span> <span class="n">fontsize</span><span class="o">=</span><span class="n">label_fontsize</span><span class="p">,</span> <span class="n">zorder</span><span class="o">=</span><span class="mi">10</span>
            <span class="p">)</span>
            <span class="n">cax_e</span><span class="o">.</span><span class="n">tick_params</span><span class="p">(</span><span class="n">axis</span><span class="o">=</span><span class="s1">&#39;both&#39;</span><span class="p">,</span> <span class="n">which</span><span class="o">=</span><span class="s1">&#39;major&#39;</span><span class="p">,</span> <span class="n">labelsize</span><span class="o">=</span><span class="n">tick_label_size</span><span class="p">)</span>

    <span class="c1">##</span>
    <span class="c1"># Draw nodes</span>
    <span class="c1">##</span>
    <span class="n">node_sizes</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="nb">len</span><span class="p">(</span><span class="n">node_rings</span><span class="p">),</span> <span class="n">N</span><span class="p">))</span>
    <span class="k">for</span> <span class="n">ring</span> <span class="ow">in</span> <span class="nb">list</span><span class="p">(</span><span class="n">node_rings</span><span class="p">):</span>  <span class="c1"># iterate through to get all node sizes</span>
        <span class="k">if</span> <span class="n">node_rings</span><span class="p">[</span><span class="n">ring</span><span class="p">][</span><span class="s2">&quot;sizes&quot;</span><span class="p">]</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">node_sizes</span><span class="p">[</span><span class="n">ring</span><span class="p">]</span> <span class="o">=</span> <span class="n">node_rings</span><span class="p">[</span><span class="n">ring</span><span class="p">][</span><span class="s2">&quot;sizes&quot;</span><span class="p">]</span>

        <span class="k">else</span><span class="p">:</span>
            <span class="n">node_sizes</span><span class="p">[</span><span class="n">ring</span><span class="p">]</span> <span class="o">=</span> <span class="n">standard_size</span>
    <span class="n">max_sizes</span> <span class="o">=</span> <span class="n">node_sizes</span><span class="o">.</span><span class="n">max</span><span class="p">(</span><span class="n">axis</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span>
    <span class="n">total_max_size</span> <span class="o">=</span> <span class="n">node_sizes</span><span class="o">.</span><span class="n">sum</span><span class="p">(</span><span class="n">axis</span><span class="o">=</span><span class="mi">0</span><span class="p">)</span><span class="o">.</span><span class="n">max</span><span class="p">()</span>
    <span class="n">node_sizes</span> <span class="o">/=</span> <span class="n">total_max_size</span>
    <span class="n">node_sizes</span> <span class="o">*=</span> <span class="n">standard_size</span>

    <span class="k">def</span> <span class="nf">get_aspect</span><span class="p">(</span><span class="n">ax</span><span class="p">):</span>
        <span class="c1"># Total figure size</span>
        <span class="n">figW</span><span class="p">,</span> <span class="n">figH</span> <span class="o">=</span> <span class="n">ax</span><span class="o">.</span><span class="n">get_figure</span><span class="p">()</span><span class="o">.</span><span class="n">get_size_inches</span><span class="p">()</span>
        <span class="c1"># print(figW, figH)</span>
        <span class="c1"># Axis size on figure</span>
        <span class="n">_</span><span class="p">,</span> <span class="n">_</span><span class="p">,</span> <span class="n">w</span><span class="p">,</span> <span class="n">h</span> <span class="o">=</span> <span class="n">ax</span><span class="o">.</span><span class="n">get_position</span><span class="p">()</span><span class="o">.</span><span class="n">bounds</span>
        <span class="c1"># Ratio of display units</span>
        <span class="c1"># print(w, h)</span>
        <span class="n">disp_ratio</span> <span class="o">=</span> <span class="p">(</span><span class="n">figH</span> <span class="o">*</span> <span class="n">h</span><span class="p">)</span> <span class="o">/</span> <span class="p">(</span><span class="n">figW</span> <span class="o">*</span> <span class="n">w</span><span class="p">)</span>
        <span class="c1"># Ratio of data units</span>
        <span class="c1"># Negative over negative because of the order of subtraction</span>
        <span class="n">data_ratio</span> <span class="o">=</span> <span class="n">sub</span><span class="p">(</span><span class="o">*</span><span class="n">ax</span><span class="o">.</span><span class="n">get_ylim</span><span class="p">())</span> <span class="o">/</span> <span class="n">sub</span><span class="p">(</span><span class="o">*</span><span class="n">ax</span><span class="o">.</span><span class="n">get_xlim</span><span class="p">())</span>
        <span class="c1"># print(data_ratio, disp_ratio)</span>
        <span class="k">return</span> <span class="n">disp_ratio</span> <span class="o">/</span> <span class="n">data_ratio</span>

    <span class="k">if</span> <span class="n">node_aspect</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">node_aspect</span> <span class="o">=</span> <span class="n">get_aspect</span><span class="p">(</span><span class="n">ax</span><span class="p">)</span>

    <span class="c1"># start drawing the outer ring first...</span>
    <span class="k">for</span> <span class="n">ring</span> <span class="ow">in</span> <span class="nb">list</span><span class="p">(</span><span class="n">node_rings</span><span class="p">)[::</span><span class="o">-</span><span class="mi">1</span><span class="p">]:</span>
        <span class="c1">#        print ring</span>
        <span class="c1"># dictionary of rings: {0:{&#39;sizes&#39;:(N,)-array, &#39;color_array&#39;:(N,)-array</span>
        <span class="c1"># or None, &#39;cmap&#39;:string, &#39;vmin&#39;:float or None, &#39;vmax&#39;:float or None}}</span>
        <span class="k">if</span> <span class="n">node_rings</span><span class="p">[</span><span class="n">ring</span><span class="p">][</span><span class="s2">&quot;color_array&quot;</span><span class="p">]</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">color_data</span> <span class="o">=</span> <span class="n">node_rings</span><span class="p">[</span><span class="n">ring</span><span class="p">][</span><span class="s2">&quot;color_array&quot;</span><span class="p">]</span>
            <span class="k">if</span> <span class="n">node_rings</span><span class="p">[</span><span class="n">ring</span><span class="p">][</span><span class="s2">&quot;vmin&quot;</span><span class="p">]</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                <span class="n">vmin</span> <span class="o">=</span> <span class="n">node_rings</span><span class="p">[</span><span class="n">ring</span><span class="p">][</span><span class="s2">&quot;vmin&quot;</span><span class="p">]</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">vmin</span> <span class="o">=</span> <span class="n">node_rings</span><span class="p">[</span><span class="n">ring</span><span class="p">][</span><span class="s2">&quot;color_array&quot;</span><span class="p">]</span><span class="o">.</span><span class="n">min</span><span class="p">()</span>
            <span class="k">if</span> <span class="n">node_rings</span><span class="p">[</span><span class="n">ring</span><span class="p">][</span><span class="s2">&quot;vmax&quot;</span><span class="p">]</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                <span class="n">vmax</span> <span class="o">=</span> <span class="n">node_rings</span><span class="p">[</span><span class="n">ring</span><span class="p">][</span><span class="s2">&quot;vmax&quot;</span><span class="p">]</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">vmax</span> <span class="o">=</span> <span class="n">node_rings</span><span class="p">[</span><span class="n">ring</span><span class="p">][</span><span class="s2">&quot;color_array&quot;</span><span class="p">]</span><span class="o">.</span><span class="n">max</span><span class="p">()</span>
            <span class="k">if</span> <span class="n">node_rings</span><span class="p">[</span><span class="n">ring</span><span class="p">][</span><span class="s2">&quot;cmap&quot;</span><span class="p">]</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                <span class="n">cmap</span> <span class="o">=</span> <span class="n">node_rings</span><span class="p">[</span><span class="n">ring</span><span class="p">][</span><span class="s2">&quot;cmap&quot;</span><span class="p">]</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">cmap</span> <span class="o">=</span> <span class="n">standard_cmap</span>
            <span class="n">data_to_rgb</span> <span class="o">=</span> <span class="n">pyplot</span><span class="o">.</span><span class="n">cm</span><span class="o">.</span><span class="n">ScalarMappable</span><span class="p">(</span>
                <span class="n">norm</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">cmap</span><span class="o">=</span><span class="n">pyplot</span><span class="o">.</span><span class="n">get_cmap</span><span class="p">(</span><span class="n">cmap</span><span class="p">)</span>
            <span class="p">)</span>
            <span class="n">data_to_rgb</span><span class="o">.</span><span class="n">set_array</span><span class="p">(</span><span class="n">color_data</span><span class="p">)</span>
            <span class="n">data_to_rgb</span><span class="o">.</span><span class="n">set_clim</span><span class="p">(</span><span class="n">vmin</span><span class="o">=</span><span class="n">vmin</span><span class="p">,</span> <span class="n">vmax</span><span class="o">=</span><span class="n">vmax</span><span class="p">)</span>
            <span class="n">colors</span> <span class="o">=</span> <span class="p">[</span><span class="n">data_to_rgb</span><span class="o">.</span><span class="n">to_rgba</span><span class="p">(</span><span class="n">color_data</span><span class="p">[</span><span class="n">n</span><span class="p">])</span> <span class="k">for</span> <span class="n">n</span> <span class="ow">in</span> <span class="n">G</span><span class="p">]</span>

            <span class="k">if</span> <span class="n">node_rings</span><span class="p">[</span><span class="n">ring</span><span class="p">][</span><span class="s2">&quot;colorbar&quot;</span><span class="p">]:</span>
                <span class="c1"># Create colorbars for nodes</span>
                <span class="c1"># cax_n = pyplot.axes([.8 + ring*0.11,</span>
                <span class="c1"># ax.get_subplotspec().get_position(ax.figure).bounds[1]+0.05, 0.025, 0.35], frameon=False) #</span>
                <span class="c1"># setup colorbar axes.</span>
                <span class="c1"># setup colorbar axes.</span>
                <span class="n">bbox_ax</span> <span class="o">=</span> <span class="n">ax</span><span class="o">.</span><span class="n">get_position</span><span class="p">()</span>
                <span class="c1"># print(bbox_ax.xmin, bbox_ax.xmax, bbox_ax.ymin, bbox_ax.ymax) </span>
                <span class="n">cax_n</span> <span class="o">=</span> <span class="n">ax</span><span class="o">.</span><span class="n">inset_axes</span><span class="p">(</span>
                    <span class="p">[</span>
                    <span class="mf">0.05</span><span class="p">,</span> <span class="o">-</span><span class="mf">0.07</span><span class="p">,</span> <span class="mf">0.4</span><span class="p">,</span> <span class="mf">0.07</span>
                        <span class="c1"># bbox_ax.xmin + width*0.05,</span>
                        <span class="c1"># bbox_ax.ymin-0.075*height+network_lower_bound-0.15,</span>
                        <span class="c1"># width*0.4,</span>
                        <span class="c1"># 0.075*height,   #0.025 + (len(all_links_edge_weights) == 0) * 0.035,</span>
                    <span class="p">],</span>
                    <span class="n">frameon</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
                <span class="p">)</span>
                <span class="n">cb_n</span> <span class="o">=</span> <span class="n">pyplot</span><span class="o">.</span><span class="n">colorbar</span><span class="p">(</span><span class="n">data_to_rgb</span><span class="p">,</span> <span class="n">cax</span><span class="o">=</span><span class="n">cax_n</span><span class="p">,</span> <span class="n">orientation</span><span class="o">=</span><span class="s2">&quot;horizontal&quot;</span><span class="p">)</span>
                <span class="c1"># try:</span>
                <span class="n">ticks_here</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span>
                    <span class="n">_myround</span><span class="p">(</span><span class="n">vmin</span><span class="p">,</span> <span class="n">node_rings</span><span class="p">[</span><span class="n">ring</span><span class="p">][</span><span class="s2">&quot;ticks&quot;</span><span class="p">],</span> <span class="s2">&quot;down&quot;</span><span class="p">),</span>
                    <span class="n">_myround</span><span class="p">(</span><span class="n">vmax</span><span class="p">,</span> <span class="n">node_rings</span><span class="p">[</span><span class="n">ring</span><span class="p">][</span><span class="s2">&quot;ticks&quot;</span><span class="p">],</span> <span class="s2">&quot;up&quot;</span><span class="p">)</span>
                    <span class="o">+</span> <span class="n">node_rings</span><span class="p">[</span><span class="n">ring</span><span class="p">][</span><span class="s2">&quot;ticks&quot;</span><span class="p">],</span>
                    <span class="n">node_rings</span><span class="p">[</span><span class="n">ring</span><span class="p">][</span><span class="s2">&quot;ticks&quot;</span><span class="p">],</span>
                <span class="p">)</span>
                <span class="n">cb_n</span><span class="o">.</span><span class="n">set_ticks</span><span class="p">(</span><span class="n">ticks_here</span><span class="p">[(</span><span class="n">vmin</span> <span class="o">&lt;=</span> <span class="n">ticks_here</span><span class="p">)</span> <span class="o">&amp;</span> <span class="p">(</span><span class="n">ticks_here</span> <span class="o">&lt;=</span> <span class="n">vmax</span><span class="p">)])</span>
                <span class="c1"># except:</span>
                <span class="c1">#     print (&#39;no ticks given&#39;)</span>
                <span class="n">cb_n</span><span class="o">.</span><span class="n">outline</span><span class="o">.</span><span class="n">clear</span><span class="p">()</span>
                <span class="c1"># cb_n.set_ticks()</span>
                <span class="n">cax_n</span><span class="o">.</span><span class="n">set_xlabel</span><span class="p">(</span>
                    <span class="n">node_rings</span><span class="p">[</span><span class="n">ring</span><span class="p">][</span><span class="s2">&quot;label&quot;</span><span class="p">],</span> <span class="n">labelpad</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span> <span class="n">fontsize</span><span class="o">=</span><span class="n">label_fontsize</span>
                <span class="p">)</span>
                <span class="n">cax_n</span><span class="o">.</span><span class="n">tick_params</span><span class="p">(</span><span class="n">axis</span><span class="o">=</span><span class="s1">&#39;both&#39;</span><span class="p">,</span> <span class="n">which</span><span class="o">=</span><span class="s1">&#39;major&#39;</span><span class="p">,</span> <span class="n">labelsize</span><span class="o">=</span><span class="n">tick_label_size</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">colors</span> <span class="o">=</span> <span class="kc">None</span>
            <span class="n">vmin</span> <span class="o">=</span> <span class="kc">None</span>
            <span class="n">vmax</span> <span class="o">=</span> <span class="kc">None</span>

        <span class="k">for</span> <span class="n">n</span> <span class="ow">in</span> <span class="n">G</span><span class="p">:</span>
            <span class="k">if</span> <span class="nb">type</span><span class="p">(</span><span class="n">node_alpha</span><span class="p">)</span> <span class="o">==</span> <span class="nb">dict</span><span class="p">:</span>
                <span class="n">alpha</span> <span class="o">=</span> <span class="n">node_alpha</span><span class="p">[</span><span class="n">n</span><span class="p">]</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">alpha</span> <span class="o">=</span> <span class="mf">1.0</span>

            <span class="k">if</span> <span class="n">special_nodes</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                <span class="k">if</span> <span class="n">n</span> <span class="ow">in</span> <span class="n">special_nodes</span><span class="p">:</span>
                    <span class="n">color_here</span> <span class="o">=</span> <span class="n">special_nodes</span><span class="p">[</span><span class="n">n</span><span class="p">]</span>
                <span class="k">else</span><span class="p">:</span>
                    <span class="n">color_here</span> <span class="o">=</span> <span class="s1">&#39;grey&#39;</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="k">if</span> <span class="n">colors</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
                    <span class="n">color_here</span> <span class="o">=</span> <span class="n">standard_color_nodes</span>
                <span class="k">else</span><span class="p">:</span>
                    <span class="n">color_here</span> <span class="o">=</span> <span class="n">colors</span><span class="p">[</span><span class="n">n</span><span class="p">]</span>

            <span class="n">c</span> <span class="o">=</span> <span class="n">Ellipse</span><span class="p">(</span>
                <span class="n">pos</span><span class="p">[</span><span class="n">n</span><span class="p">],</span>
                <span class="n">width</span><span class="o">=</span><span class="n">node_sizes</span><span class="p">[:</span> <span class="n">ring</span> <span class="o">+</span> <span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">sum</span><span class="p">(</span><span class="n">axis</span><span class="o">=</span><span class="mi">0</span><span class="p">)[</span><span class="n">n</span><span class="p">]</span> <span class="o">*</span> <span class="n">node_aspect</span><span class="p">,</span>
                <span class="n">height</span><span class="o">=</span><span class="n">node_sizes</span><span class="p">[:</span> <span class="n">ring</span> <span class="o">+</span> <span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">sum</span><span class="p">(</span><span class="n">axis</span><span class="o">=</span><span class="mi">0</span><span class="p">)[</span><span class="n">n</span><span class="p">],</span>
                <span class="n">clip_on</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
                <span class="n">facecolor</span><span class="o">=</span><span class="n">color_here</span><span class="p">,</span>
                <span class="n">edgecolor</span><span class="o">=</span><span class="n">color_here</span><span class="p">,</span>
                <span class="n">zorder</span><span class="o">=-</span><span class="n">ring</span> <span class="o">-</span> <span class="mi">1</span> <span class="o">+</span> <span class="mi">2</span><span class="p">,</span>
                <span class="n">transform</span><span class="o">=</span><span class="n">transform</span><span class="p">,</span>
            <span class="p">)</span>

            <span class="c1"># else:</span>
            <span class="c1">#     if special_nodes is not None and n in special_nodes:</span>
            <span class="c1">#         color_here = special_nodes[n]</span>
            <span class="c1">#     else:</span>
            <span class="c1">#         color_here = colors[n]</span>
            <span class="c1">#     c = Ellipse(</span>
            <span class="c1">#         pos[n],</span>
            <span class="c1">#         width=node_sizes[: ring + 1].sum(axis=0)[n] * node_aspect,</span>
            <span class="c1">#         height=node_sizes[: ring + 1].sum(axis=0)[n],</span>
            <span class="c1">#         clip_on=False,</span>
            <span class="c1">#         facecolor=colors[n],</span>
            <span class="c1">#         edgecolor=colors[n],</span>
            <span class="c1">#         zorder=-ring - 1,</span>
            <span class="c1">#     )</span>

            <span class="n">ax</span><span class="o">.</span><span class="n">add_patch</span><span class="p">(</span><span class="n">c</span><span class="p">)</span>

            <span class="k">if</span> <span class="n">node_classification</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> <span class="n">node_classification</span><span class="p">[</span><span class="n">n</span><span class="p">]</span> <span class="ow">in</span> <span class="p">[</span><span class="s2">&quot;space_context_last&quot;</span><span class="p">,</span> <span class="s2">&quot;space_dummy_last&quot;</span><span class="p">,</span> <span class="s2">&quot;time_dummy_last&quot;</span><span class="p">]:</span>
                <span class="n">node_height</span> <span class="o">=</span> <span class="n">node_sizes</span><span class="p">[:</span> <span class="n">ring</span> <span class="o">+</span> <span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">sum</span><span class="p">(</span><span class="n">axis</span><span class="o">=</span><span class="mi">0</span><span class="p">)[</span><span class="n">n</span><span class="p">]</span>
                <span class="n">node_width_difference_to_height</span> <span class="o">=</span> <span class="n">node_height</span> <span class="o">*</span> <span class="p">(</span><span class="mi">1</span> <span class="o">-</span> <span class="n">node_aspect</span><span class="p">)</span>

                <span class="n">c_wide</span> <span class="o">=</span> <span class="n">mpatches</span><span class="o">.</span><span class="n">FancyBboxPatch</span><span class="p">((</span><span class="n">pos</span><span class="p">[</span><span class="n">n</span><span class="o">-</span><span class="n">max_lag</span><span class="o">+</span><span class="mi">1</span><span class="p">][</span><span class="mi">0</span><span class="p">]</span> <span class="o">+</span> <span class="n">node_width_difference_to_height</span> <span class="o">/</span> <span class="mi">2</span><span class="p">,</span> <span class="n">pos</span><span class="p">[</span><span class="n">n</span><span class="o">-</span><span class="n">max_lag</span><span class="o">+</span><span class="mi">1</span><span class="p">][</span><span class="mi">1</span><span class="p">]),</span>
                                                 <span class="p">(</span><span class="n">pos</span><span class="p">[</span><span class="n">n</span><span class="p">][</span><span class="mi">0</span><span class="p">]</span> <span class="o">-</span> <span class="n">pos</span><span class="p">[</span><span class="n">n</span><span class="o">-</span><span class="n">max_lag</span><span class="o">+</span><span class="mi">1</span><span class="p">][</span><span class="mi">0</span><span class="p">]</span> <span class="o">-</span> <span class="n">node_width_difference_to_height</span><span class="p">),</span>
                                                 <span class="mf">0.</span><span class="p">,</span>
                                                 <span class="n">boxstyle</span><span class="o">=</span><span class="n">mpatches</span><span class="o">.</span><span class="n">BoxStyle</span><span class="o">.</span><span class="n">Round</span><span class="p">(</span><span class="n">pad</span><span class="o">=</span><span class="mf">0.5</span> <span class="o">*</span> <span class="n">node_height</span><span class="p">),</span>
                                                 <span class="n">facecolor</span><span class="o">=</span><span class="n">color_here</span><span class="p">,</span>
                                                 <span class="n">edgecolor</span><span class="o">=</span><span class="n">color_here</span><span class="p">,</span>
                                                 <span class="p">)</span>

                <span class="n">ax</span><span class="o">.</span><span class="n">add_patch</span><span class="p">(</span><span class="n">c_wide</span><span class="p">)</span>


            <span class="c1"># avoiding attribute error raised by changes in networkx</span>
            <span class="k">if</span> <span class="nb">hasattr</span><span class="p">(</span><span class="n">G</span><span class="p">,</span> <span class="s2">&quot;node&quot;</span><span class="p">):</span>
                <span class="c1"># works with networkx 1.10</span>
                <span class="n">G</span><span class="o">.</span><span class="n">node</span><span class="p">[</span><span class="n">n</span><span class="p">][</span><span class="s2">&quot;patch&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">c</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="c1"># works with networkx 2.4</span>
                <span class="n">G</span><span class="o">.</span><span class="n">nodes</span><span class="p">[</span><span class="n">n</span><span class="p">][</span><span class="s2">&quot;patch&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">c</span>

            <span class="k">if</span> <span class="n">ring</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
                <span class="n">ax</span><span class="o">.</span><span class="n">text</span><span class="p">(</span>
                    <span class="n">pos</span><span class="p">[</span><span class="n">n</span><span class="p">][</span><span class="mi">0</span><span class="p">],</span>
                    <span class="n">pos</span><span class="p">[</span><span class="n">n</span><span class="p">][</span><span class="mi">1</span><span class="p">],</span>
                    <span class="n">node_labels</span><span class="p">[</span><span class="n">n</span><span class="p">],</span>
                    <span class="n">fontsize</span><span class="o">=</span><span class="n">node_label_size</span><span class="p">,</span>
                    <span class="n">horizontalalignment</span><span class="o">=</span><span class="s2">&quot;center&quot;</span><span class="p">,</span>
                    <span class="n">verticalalignment</span><span class="o">=</span><span class="s2">&quot;center&quot;</span><span class="p">,</span>
                    <span class="n">alpha</span><span class="o">=</span><span class="mf">1.0</span><span class="p">,</span>
                    <span class="n">zorder</span><span class="o">=</span><span class="mf">5.</span><span class="p">,</span>
                    <span class="n">transform</span><span class="o">=</span><span class="n">transform</span><span class="p">,</span>
                <span class="p">)</span>
                <span class="k">if</span> <span class="n">show_autodependency_lags</span><span class="p">:</span>
                    <span class="n">ax</span><span class="o">.</span><span class="n">text</span><span class="p">(</span>
                        <span class="n">pos</span><span class="p">[</span><span class="n">n</span><span class="p">][</span><span class="mi">0</span><span class="p">],</span>
                        <span class="n">pos</span><span class="p">[</span><span class="n">n</span><span class="p">][</span><span class="mi">1</span><span class="p">],</span>
                        <span class="n">autodep_sig_lags</span><span class="p">[</span><span class="n">n</span><span class="p">],</span>
                        <span class="n">fontsize</span><span class="o">=</span><span class="n">link_label_fontsize</span><span class="p">,</span>
                        <span class="n">horizontalalignment</span><span class="o">=</span><span class="s2">&quot;center&quot;</span><span class="p">,</span>
                        <span class="n">verticalalignment</span><span class="o">=</span><span class="s2">&quot;center&quot;</span><span class="p">,</span>
                        <span class="n">color</span><span class="o">=</span><span class="s2">&quot;black&quot;</span><span class="p">,</span>
                        <span class="n">zorder</span><span class="o">=</span><span class="mf">5.</span><span class="p">,</span>
                        <span class="n">transform</span><span class="o">=</span><span class="n">transform</span><span class="p">,</span>
                    <span class="p">)</span>

    <span class="c1"># Draw edges</span>
    <span class="n">seen</span> <span class="o">=</span> <span class="p">{}</span>
    <span class="k">for</span> <span class="p">(</span><span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">,</span> <span class="n">d</span><span class="p">)</span> <span class="ow">in</span> <span class="n">G</span><span class="o">.</span><span class="n">edges</span><span class="p">(</span><span class="n">data</span><span class="o">=</span><span class="kc">True</span><span class="p">):</span>
        <span class="k">if</span> <span class="n">d</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;no_links&quot;</span><span class="p">):</span>
            <span class="n">d</span><span class="p">[</span><span class="s2">&quot;inner_edge_alpha&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="mf">1e-8</span>
            <span class="n">d</span><span class="p">[</span><span class="s2">&quot;outer_edge_alpha&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="mf">1e-8</span>
        <span class="k">if</span> <span class="n">u</span> <span class="o">!=</span> <span class="n">v</span><span class="p">:</span>
            <span class="k">if</span> <span class="n">d</span><span class="p">[</span><span class="s2">&quot;outer_edge&quot;</span><span class="p">]:</span>
                <span class="n">seen</span><span class="p">[(</span><span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">)]</span> <span class="o">=</span> <span class="n">draw_edge</span><span class="p">(</span><span class="n">ax</span><span class="p">,</span> <span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">,</span> <span class="n">d</span><span class="p">,</span> <span class="n">seen</span><span class="p">,</span> <span class="n">outer_edge</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
            <span class="k">if</span> <span class="n">d</span><span class="p">[</span><span class="s2">&quot;inner_edge&quot;</span><span class="p">]:</span>
                <span class="n">seen</span><span class="p">[(</span><span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">)]</span> <span class="o">=</span> <span class="n">draw_edge</span><span class="p">(</span><span class="n">ax</span><span class="p">,</span> <span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">,</span> <span class="n">d</span><span class="p">,</span> <span class="n">seen</span><span class="p">,</span> <span class="n">outer_edge</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span>

    <span class="c1"># if network_left_bound is not None:</span>
    <span class="c1">#     network_right_bound = 0.98</span>
    <span class="c1"># else:</span>
    <span class="c1">#     network_right_bound = None</span>
    <span class="c1"># fig.subplots_adjust(bottom=network_lower_bound, left=network_left_bound, right=network_right_bound) #, right=0.97)</span>


<div class="viewcode-block" id="plot_graph"><a class="viewcode-back" href="../../index.html#tigramite.plotting.plot_graph">[docs]</a><span class="k">def</span> <span class="nf">plot_graph</span><span class="p">(</span>
    <span class="n">graph</span><span class="p">,</span>
    <span class="n">val_matrix</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
    <span class="n">var_names</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
    <span class="n">fig_ax</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
    <span class="n">figsize</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
    <span class="n">save_name</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
    <span class="n">link_colorbar_label</span><span class="o">=</span><span class="s2">&quot;MCI&quot;</span><span class="p">,</span>
    <span class="n">node_colorbar_label</span><span class="o">=</span><span class="s2">&quot;auto-MCI&quot;</span><span class="p">,</span>
    <span class="n">link_width</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
    <span class="n">link_attribute</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
    <span class="n">node_pos</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
    <span class="n">arrow_linewidth</span><span class="o">=</span><span class="mf">8.0</span><span class="p">,</span>
    <span class="n">vmin_edges</span><span class="o">=-</span><span class="mi">1</span><span class="p">,</span>
    <span class="n">vmax_edges</span><span class="o">=</span><span class="mf">1.0</span><span class="p">,</span>
    <span class="n">edge_ticks</span><span class="o">=</span><span class="mf">0.4</span><span class="p">,</span>
    <span class="n">cmap_edges</span><span class="o">=</span><span class="s2">&quot;RdBu_r&quot;</span><span class="p">,</span>
    <span class="n">vmin_nodes</span><span class="o">=-</span><span class="mi">1</span><span class="p">,</span>
    <span class="n">vmax_nodes</span><span class="o">=</span><span class="mf">1.0</span><span class="p">,</span>
    <span class="n">node_ticks</span><span class="o">=</span><span class="mf">0.4</span><span class="p">,</span>
    <span class="n">cmap_nodes</span><span class="o">=</span><span class="s2">&quot;RdBu_r&quot;</span><span class="p">,</span>
    <span class="n">node_size</span><span class="o">=</span><span class="mf">0.3</span><span class="p">,</span>
    <span class="n">node_aspect</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
    <span class="n">arrowhead_size</span><span class="o">=</span><span class="mi">20</span><span class="p">,</span>
    <span class="n">curved_radius</span><span class="o">=</span><span class="mf">0.2</span><span class="p">,</span>
    <span class="n">label_fontsize</span><span class="o">=</span><span class="mi">10</span><span class="p">,</span>
    <span class="n">tick_label_size</span><span class="o">=</span><span class="mi">6</span><span class="p">,</span>
    <span class="n">alpha</span><span class="o">=</span><span class="mf">1.0</span><span class="p">,</span>
    <span class="n">node_label_size</span><span class="o">=</span><span class="mi">10</span><span class="p">,</span>
    <span class="n">link_label_fontsize</span><span class="o">=</span><span class="mi">10</span><span class="p">,</span>
    <span class="n">lag_array</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
    <span class="c1"># network_lower_bound=0.2,</span>
    <span class="n">show_colorbar</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
    <span class="n">inner_edge_style</span><span class="o">=</span><span class="s2">&quot;dashed&quot;</span><span class="p">,</span>
    <span class="n">link_matrix</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
    <span class="n">special_nodes</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
    <span class="n">show_autodependency_lags</span><span class="o">=</span><span class="kc">False</span>
<span class="p">):</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;Creates a network plot.</span>
<span class="sd">    </span>
<span class="sd">    This is still in beta. The network is defined from links in graph. Nodes</span>
<span class="sd">    denote variables, straight links contemporaneous dependencies and curved</span>
<span class="sd">    arrows lagged dependencies. The node color denotes the maximal absolute</span>
<span class="sd">    auto-dependency and the link color the value at the lag with maximal</span>
<span class="sd">    absolute cross-dependency. The link label lists the lags with significant</span>
<span class="sd">    dependency in order of absolute magnitude. The network can also be</span>
<span class="sd">    plotted over a map drawn before on the same axis. Then the node positions</span>
<span class="sd">    can be supplied in appropriate axis coordinates via node_pos.</span>

<span class="sd">    Parameters</span>
<span class="sd">    ----------</span>
<span class="sd">    graph : string or bool array-like, optional (default: None)</span>
<span class="sd">        Either string matrix providing graph or bool array providing only adjacencies</span>
<span class="sd">        Must be of same shape as val_matrix. </span>
<span class="sd">    val_matrix : array_like</span>
<span class="sd">        Matrix of shape (N, N, tau_max+1) containing test statistic values.</span>
<span class="sd">    var_names : list, optional (default: None)</span>
<span class="sd">        List of variable names. If None, range(N) is used.</span>
<span class="sd">    fig_ax : tuple of figure and axis object, optional (default: None)</span>
<span class="sd">        Figure and axes instance. If None they are created.</span>
<span class="sd">    figsize : tuple</span>
<span class="sd">        Size of figure.</span>
<span class="sd">    save_name : str, optional (default: None)</span>
<span class="sd">        Name of figure file to save figure. If None, figure is shown in window.</span>
<span class="sd">    link_colorbar_label : str, optional (default: &#39;MCI&#39;)</span>
<span class="sd">        Test statistic label.</span>
<span class="sd">    node_colorbar_label : str, optional (default: &#39;auto-MCI&#39;)</span>
<span class="sd">        Test statistic label for auto-dependencies.</span>
<span class="sd">    link_width : array-like, optional (default: None)</span>
<span class="sd">        Array of val_matrix.shape specifying relative link width with maximum</span>
<span class="sd">        given by arrow_linewidth. If None, all links have same width.</span>
<span class="sd">    link_attribute : array-like, optional (default: None)</span>
<span class="sd">        String array of val_matrix.shape specifying link attributes.</span>
<span class="sd">    node_pos : dictionary, optional (default: None)</span>
<span class="sd">        Dictionary of node positions in axis coordinates of form</span>
<span class="sd">        node_pos = {&#39;x&#39;:array of shape (N,), &#39;y&#39;:array of shape(N)}. These</span>
<span class="sd">        coordinates could have been transformed before for basemap plots. You can</span>
<span class="sd">        also add a key &#39;transform&#39;:ccrs.PlateCarree() in order to plot graphs on </span>
<span class="sd">        a map using cartopy.</span>
<span class="sd">    arrow_linewidth : float, optional (default: 30)</span>
<span class="sd">        Linewidth.</span>
<span class="sd">    vmin_edges : float, optional (default: -1)</span>
<span class="sd">        Link colorbar scale lower bound.</span>
<span class="sd">    vmax_edges : float, optional (default: 1)</span>
<span class="sd">        Link colorbar scale upper bound.</span>
<span class="sd">    edge_ticks : float, optional (default: 0.4)</span>
<span class="sd">        Link tick mark interval.</span>
<span class="sd">    cmap_edges : str, optional (default: &#39;RdBu_r&#39;)</span>
<span class="sd">        Colormap for links.</span>
<span class="sd">    vmin_nodes : float, optional (default: 0)</span>
<span class="sd">        Node colorbar scale lower bound.</span>
<span class="sd">    vmax_nodes : float, optional (default: 1)</span>
<span class="sd">        Node colorbar scale upper bound.</span>
<span class="sd">    node_ticks : float, optional (default: 0.4)</span>
<span class="sd">        Node tick mark interval.</span>
<span class="sd">    cmap_nodes : str, optional (default: &#39;OrRd&#39;)</span>
<span class="sd">        Colormap for links.</span>
<span class="sd">    node_size : int, optional (default: 0.3)</span>
<span class="sd">        Node size.</span>
<span class="sd">    node_aspect : float, optional (default: None)</span>
<span class="sd">        Ratio between the heigth and width of the varible nodes.</span>
<span class="sd">    arrowhead_size : int, optional (default: 20)</span>
<span class="sd">        Size of link arrow head. Passed on to FancyArrowPatch object.</span>
<span class="sd">    curved_radius, float, optional (default: 0.2)</span>
<span class="sd">        Curvature of links. Passed on to FancyArrowPatch object.</span>
<span class="sd">    label_fontsize : int, optional (default: 10)</span>
<span class="sd">        Fontsize of colorbar labels.</span>
<span class="sd">    alpha : float, optional (default: 1.)</span>
<span class="sd">        Opacity.</span>
<span class="sd">    node_label_size : int, optional (default: 10)</span>
<span class="sd">        Fontsize of node labels.</span>
<span class="sd">    link_label_fontsize : int, optional (default: 6)</span>
<span class="sd">        Fontsize of link labels.</span>
<span class="sd">    tick_label_size : int, optional (default: 6)</span>
<span class="sd">        Fontsize of tick labels.</span>
<span class="sd">    lag_array : array, optional (default: None)</span>
<span class="sd">        Optional specification of lags overwriting np.arange(0, tau_max+1)</span>
<span class="sd">    show_colorbar : bool</span>
<span class="sd">        Whether to show colorbars for links and nodes.</span>
<span class="sd">    show_autodependency_lags : bool (default: False)</span>
<span class="sd">        Shows significant autodependencies for a node.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">if</span> <span class="n">link_matrix</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;link_matrix is deprecated and replaced by graph array&quot;</span>
                         <span class="s2">&quot; which is now returned by all methods.&quot;</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">fig_ax</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">fig</span> <span class="o">=</span> <span class="n">pyplot</span><span class="o">.</span><span class="n">figure</span><span class="p">(</span><span class="n">figsize</span><span class="o">=</span><span class="n">figsize</span><span class="p">)</span>
        <span class="n">ax</span> <span class="o">=</span> <span class="n">fig</span><span class="o">.</span><span class="n">add_subplot</span><span class="p">(</span><span class="mi">111</span><span class="p">,</span> <span class="n">frame_on</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">fig</span><span class="p">,</span> <span class="n">ax</span> <span class="o">=</span> <span class="n">fig_ax</span>

    <span class="n">graph</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">copy</span><span class="p">(</span><span class="n">graph</span><span class="o">.</span><span class="n">squeeze</span><span class="p">())</span>

    <span class="k">if</span> <span class="n">graph</span><span class="o">.</span><span class="n">ndim</span> <span class="o">==</span> <span class="mi">4</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;Time series graph of shape (N,N,tau_max+1,tau_max+1) cannot be represented by plot_graph,&quot;</span>
                         <span class="s2">&quot; use plot_time_series_graph instead.&quot;</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">graph</span><span class="o">.</span><span class="n">ndim</span> <span class="o">==</span> <span class="mi">2</span><span class="p">:</span>
        <span class="c1"># If a non-time series (N,N)-graph is given, insert a dummy dimension</span>
        <span class="n">graph</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">expand_dims</span><span class="p">(</span><span class="n">graph</span><span class="p">,</span> <span class="n">axis</span> <span class="o">=</span> <span class="mi">2</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">val_matrix</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">no_coloring</span> <span class="o">=</span> <span class="kc">True</span>
        <span class="n">cmap_edges</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="n">cmap_nodes</span> <span class="o">=</span> <span class="kc">None</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">no_coloring</span> <span class="o">=</span> <span class="kc">False</span>

    <span class="p">(</span><span class="n">graph</span><span class="p">,</span> <span class="n">val_matrix</span><span class="p">,</span> <span class="n">link_width</span><span class="p">,</span> <span class="n">link_attribute</span><span class="p">)</span> <span class="o">=</span> <span class="n">_check_matrices</span><span class="p">(</span>
        <span class="n">graph</span><span class="p">,</span> <span class="n">val_matrix</span><span class="p">,</span> <span class="n">link_width</span><span class="p">,</span> <span class="n">link_attribute</span><span class="p">)</span>
    

    <span class="n">N</span><span class="p">,</span> <span class="n">N</span><span class="p">,</span> <span class="n">dummy</span> <span class="o">=</span> <span class="n">graph</span><span class="o">.</span><span class="n">shape</span>
    <span class="n">tau_max</span> <span class="o">=</span> <span class="n">dummy</span> <span class="o">-</span> <span class="mi">1</span>
    <span class="n">max_lag</span> <span class="o">=</span> <span class="n">tau_max</span> <span class="o">+</span> <span class="mi">1</span>

    <span class="k">if</span> <span class="n">np</span><span class="o">.</span><span class="n">count_nonzero</span><span class="p">(</span><span class="n">graph</span> <span class="o">!=</span> <span class="s2">&quot;&quot;</span><span class="p">)</span> <span class="o">==</span> <span class="n">np</span><span class="o">.</span><span class="n">count_nonzero</span><span class="p">(</span>
        <span class="n">np</span><span class="o">.</span><span class="n">diagonal</span><span class="p">(</span><span class="n">graph</span><span class="p">)</span> <span class="o">!=</span> <span class="s2">&quot;&quot;</span>
    <span class="p">):</span>
        <span class="n">diagonal</span> <span class="o">=</span> <span class="kc">True</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">diagonal</span> <span class="o">=</span> <span class="kc">False</span>

    <span class="k">if</span> <span class="n">np</span><span class="o">.</span><span class="n">count_nonzero</span><span class="p">(</span><span class="n">graph</span> <span class="o">==</span> <span class="s2">&quot;&quot;</span><span class="p">)</span> <span class="o">==</span> <span class="n">graph</span><span class="o">.</span><span class="n">size</span> <span class="ow">or</span> <span class="n">diagonal</span><span class="p">:</span>
        <span class="n">graph</span><span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">0</span><span class="p">]</span> <span class="o">=</span> <span class="s2">&quot;xxx&quot;</span>  <span class="c1"># Workaround, will not be plotted... </span>
        <span class="n">no_links</span> <span class="o">=</span> <span class="kc">True</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">no_links</span> <span class="o">=</span> <span class="kc">False</span>

    <span class="k">if</span> <span class="n">var_names</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">var_names</span> <span class="o">=</span> <span class="nb">range</span><span class="p">(</span><span class="n">N</span><span class="p">)</span>

    <span class="c1"># Define graph links by absolute maximum (positive or negative like for</span>
    <span class="c1"># partial correlation)</span>
    <span class="c1"># val_matrix[np.abs(val_matrix) &lt; sig_thres] = 0.</span>

    <span class="c1"># Only draw link in one direction among contemp</span>
    <span class="c1"># Remove lower triangle</span>
    <span class="n">link_matrix_upper</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">copy</span><span class="p">(</span><span class="n">graph</span><span class="p">)</span>
    <span class="n">link_matrix_upper</span><span class="p">[:,</span> <span class="p">:,</span> <span class="mi">0</span><span class="p">]</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">triu</span><span class="p">(</span><span class="n">link_matrix_upper</span><span class="p">[:,</span> <span class="p">:,</span> <span class="mi">0</span><span class="p">])</span>

    <span class="c1"># net = _get_absmax(link_matrix != &quot;&quot;)</span>
    <span class="n">net</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">any</span><span class="p">(</span><span class="n">link_matrix_upper</span> <span class="o">!=</span> <span class="s2">&quot;&quot;</span><span class="p">,</span> <span class="n">axis</span><span class="o">=</span><span class="mi">2</span><span class="p">)</span>
    <span class="n">G</span> <span class="o">=</span> <span class="n">nx</span><span class="o">.</span><span class="n">DiGraph</span><span class="p">(</span><span class="n">net</span><span class="p">)</span>
    
    <span class="c1"># This handels Graphs with no links.</span>
    <span class="c1"># nx.draw(G, alpha=0, zorder=-10)</span>

    <span class="n">node_color</span> <span class="o">=</span> <span class="nb">list</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="n">N</span><span class="p">))</span>

    <span class="k">if</span> <span class="n">show_autodependency_lags</span><span class="p">:</span>
        <span class="n">autodep_sig_lags</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">full</span><span class="p">(</span><span class="n">N</span><span class="p">,</span> <span class="kc">None</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="s1">&#39;object&#39;</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">autodep_sig_lags</span> <span class="o">=</span> <span class="kc">None</span>

    <span class="c1"># list of all strengths for color map</span>
    <span class="n">all_strengths</span> <span class="o">=</span> <span class="p">[]</span>
    <span class="c1"># Add attributes, contemporaneous and lagged links are handled separately</span>
    <span class="k">for</span> <span class="p">(</span><span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">,</span> <span class="n">dic</span><span class="p">)</span> <span class="ow">in</span> <span class="n">G</span><span class="o">.</span><span class="n">edges</span><span class="p">(</span><span class="n">data</span><span class="o">=</span><span class="kc">True</span><span class="p">):</span>
        <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;no_links&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">no_links</span>
        <span class="c1"># average lagfunc for link u --&gt; v ANDOR u -- v</span>
        <span class="k">if</span> <span class="n">tau_max</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
            <span class="c1"># argmax of absolute maximum where a link exists!</span>
            <span class="n">links</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">where</span><span class="p">(</span><span class="n">link_matrix_upper</span><span class="p">[</span><span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">,</span> <span class="mi">1</span><span class="p">:]</span> <span class="o">!=</span> <span class="s2">&quot;&quot;</span><span class="p">)[</span><span class="mi">0</span><span class="p">]</span>
            <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">links</span><span class="p">)</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
                <span class="n">argmax_links</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">abs</span><span class="p">(</span><span class="n">val_matrix</span><span class="p">[</span><span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">][</span><span class="mi">1</span><span class="p">:][</span><span class="n">links</span><span class="p">])</span><span class="o">.</span><span class="n">argmax</span><span class="p">()</span>
                <span class="n">argmax</span> <span class="o">=</span> <span class="n">links</span><span class="p">[</span><span class="n">argmax_links</span><span class="p">]</span> <span class="o">+</span> <span class="mi">1</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">argmax</span> <span class="o">=</span> <span class="mi">0</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">argmax</span> <span class="o">=</span> <span class="mi">0</span>

        <span class="k">if</span> <span class="n">u</span> <span class="o">!=</span> <span class="n">v</span><span class="p">:</span>
            <span class="c1"># For contemp links masking or finite samples can lead to different</span>
            <span class="c1"># values for u--v and v--u</span>
            <span class="c1"># Here we use the  maximum for the width and weight (=color)</span>
            <span class="c1"># of the link</span>
            <span class="c1"># Draw link if u--v OR v--u at lag 0 is nonzero</span>
            <span class="c1"># dic[&#39;inner_edge&#39;] = ((np.abs(val_matrix[u, v][0]) &gt;=</span>
            <span class="c1">#                       sig_thres[u, v][0]) or</span>
            <span class="c1">#                      (np.abs(val_matrix[v, u][0]) &gt;=</span>
            <span class="c1">#                       sig_thres[v, u][0]))</span>
            <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;inner_edge&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">link_matrix_upper</span><span class="p">[</span><span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">,</span> <span class="mi">0</span><span class="p">]</span>
            <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;inner_edge_type&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">link_matrix_upper</span><span class="p">[</span><span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">,</span> <span class="mi">0</span><span class="p">]</span>
            <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;inner_edge_alpha&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">alpha</span>
            <span class="k">if</span> <span class="n">no_coloring</span><span class="p">:</span>
                <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;inner_edge_color&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;inner_edge_color&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">val_matrix</span><span class="p">[</span><span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">,</span> <span class="mi">0</span><span class="p">]</span>
            <span class="c1"># # value at argmax of average</span>
            <span class="c1"># if np.abs(val_matrix[u, v][0] - val_matrix[v, u][0]) &gt; .0001:</span>
            <span class="c1">#     print(&quot;Contemporaneous I(%d; %d)=%.3f != I(%d; %d)=%.3f&quot; % (</span>
            <span class="c1">#           u, v, val_matrix[u, v][0], v, u, val_matrix[v, u][0]) +</span>
            <span class="c1">#           &quot; due to conditions, finite sample effects or &quot;</span>
            <span class="c1">#           &quot;masking, here edge color = &quot;</span>
            <span class="c1">#           &quot;larger (absolute) value.&quot;)</span>
            <span class="c1"># dic[&#39;inner_edge_color&#39;] = _get_absmax(</span>
            <span class="c1">#     np.array([[[val_matrix[u, v][0],</span>
            <span class="c1">#                    val_matrix[v, u][0]]]])).squeeze()</span>

            <span class="k">if</span> <span class="n">link_width</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
                <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;inner_edge_width&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">arrow_linewidth</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;inner_edge_width&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="p">(</span>
                    <span class="n">link_width</span><span class="p">[</span><span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">,</span> <span class="mi">0</span><span class="p">]</span> <span class="o">/</span> <span class="n">link_width</span><span class="o">.</span><span class="n">max</span><span class="p">()</span> <span class="o">*</span> <span class="n">arrow_linewidth</span>
                <span class="p">)</span>

            <span class="k">if</span> <span class="n">link_attribute</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
                <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;inner_edge_attribute&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;inner_edge_attribute&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">link_attribute</span><span class="p">[</span><span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">,</span> <span class="mi">0</span><span class="p">]</span>

            <span class="c1">#     # fraction of nonzero values</span>
            <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;inner_edge_style&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="s2">&quot;solid&quot;</span>
            <span class="c1"># else:</span>
            <span class="c1"># dic[&#39;inner_edge_style&#39;] = link_style[</span>
            <span class="c1">#         u, v, 0]</span>

            <span class="n">all_strengths</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">dic</span><span class="p">[</span><span class="s2">&quot;inner_edge_color&quot;</span><span class="p">])</span>

            <span class="k">if</span> <span class="n">tau_max</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
                <span class="c1"># True if ensemble mean at lags &gt; 0 is nonzero</span>
                <span class="c1"># dic[&#39;outer_edge&#39;] = np.any(</span>
                <span class="c1">#     np.abs(val_matrix[u, v][1:]) &gt;= sig_thres[u, v][1:])</span>
                <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;outer_edge&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">any</span><span class="p">(</span><span class="n">link_matrix_upper</span><span class="p">[</span><span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">,</span> <span class="mi">1</span><span class="p">:]</span> <span class="o">!=</span> <span class="s2">&quot;&quot;</span><span class="p">)</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;outer_edge&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="kc">False</span>
            <span class="c1"># print(u, v, dic[&quot;outer_edge&quot;], argmax, link_matrix_upper[u, v, :])</span>

            <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;outer_edge_type&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">link_matrix_upper</span><span class="p">[</span><span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">,</span> <span class="n">argmax</span><span class="p">]</span>

            <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;outer_edge_alpha&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">alpha</span>
            <span class="k">if</span> <span class="n">link_width</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
                <span class="c1"># fraction of nonzero values</span>
                <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;outer_edge_width&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">arrow_linewidth</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;outer_edge_width&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="p">(</span>
                    <span class="n">link_width</span><span class="p">[</span><span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">,</span> <span class="n">argmax</span><span class="p">]</span> <span class="o">/</span> <span class="n">link_width</span><span class="o">.</span><span class="n">max</span><span class="p">()</span> <span class="o">*</span> <span class="n">arrow_linewidth</span>
                <span class="p">)</span>

            <span class="k">if</span> <span class="n">link_attribute</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
                <span class="c1"># fraction of nonzero values</span>
                <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;outer_edge_attribute&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;outer_edge_attribute&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">link_attribute</span><span class="p">[</span><span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">,</span> <span class="n">argmax</span><span class="p">]</span>

            <span class="c1"># value at argmax of average</span>
            <span class="k">if</span> <span class="n">no_coloring</span><span class="p">:</span>
                <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;outer_edge_color&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;outer_edge_color&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">val_matrix</span><span class="p">[</span><span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">][</span><span class="n">argmax</span><span class="p">]</span>
            <span class="n">all_strengths</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">dic</span><span class="p">[</span><span class="s2">&quot;outer_edge_color&quot;</span><span class="p">])</span>

            <span class="c1"># Sorted list of significant lags (only if robust wrt</span>
            <span class="c1"># d[&#39;min_ensemble_frac&#39;])</span>
            <span class="k">if</span> <span class="n">tau_max</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
                <span class="n">lags</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">abs</span><span class="p">(</span><span class="n">val_matrix</span><span class="p">[</span><span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">][</span><span class="mi">1</span><span class="p">:])</span><span class="o">.</span><span class="n">argsort</span><span class="p">()[::</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span> <span class="o">+</span> <span class="mi">1</span>
                <span class="n">sig_lags</span> <span class="o">=</span> <span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">where</span><span class="p">(</span><span class="n">link_matrix_upper</span><span class="p">[</span><span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">,</span> <span class="mi">1</span><span class="p">:]</span> <span class="o">!=</span> <span class="s2">&quot;&quot;</span><span class="p">)[</span><span class="mi">0</span><span class="p">]</span> <span class="o">+</span> <span class="mi">1</span><span class="p">)</span><span class="o">.</span><span class="n">tolist</span><span class="p">()</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">lags</span><span class="p">,</span> <span class="n">sig_lags</span> <span class="o">=</span> <span class="p">[],</span> <span class="p">[]</span>
            <span class="k">if</span> <span class="n">lag_array</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;label&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="nb">str</span><span class="p">([</span><span class="n">lag_array</span><span class="p">[</span><span class="n">l</span><span class="p">]</span> <span class="k">for</span> <span class="n">l</span> <span class="ow">in</span> <span class="n">lags</span> <span class="k">if</span> <span class="n">l</span> <span class="ow">in</span> <span class="n">sig_lags</span><span class="p">])[</span><span class="mi">1</span><span class="p">:</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">replace</span><span class="p">(</span><span class="s2">&quot; &quot;</span><span class="p">,</span> <span class="s2">&quot;&quot;</span><span class="p">)</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;label&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="nb">str</span><span class="p">([</span><span class="n">l</span> <span class="k">for</span> <span class="n">l</span> <span class="ow">in</span> <span class="n">lags</span> <span class="k">if</span> <span class="n">l</span> <span class="ow">in</span> <span class="n">sig_lags</span><span class="p">])[</span><span class="mi">1</span><span class="p">:</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">replace</span><span class="p">(</span><span class="s2">&quot; &quot;</span><span class="p">,</span> <span class="s2">&quot;&quot;</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="c1"># Node color is max of average autodependency</span>
            <span class="k">if</span> <span class="n">no_coloring</span><span class="p">:</span>
                <span class="n">node_color</span><span class="p">[</span><span class="n">u</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">node_color</span><span class="p">[</span><span class="n">u</span><span class="p">]</span> <span class="o">=</span> <span class="n">val_matrix</span><span class="p">[</span><span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">][</span><span class="n">argmax</span><span class="p">]</span>

            <span class="k">if</span> <span class="n">show_autodependency_lags</span><span class="p">:</span>
                <span class="n">autodep_sig_lags</span><span class="p">[</span><span class="n">u</span><span class="p">]</span> <span class="o">=</span> <span class="s2">&quot;</span><span class="se">\n\n\n</span><span class="s2">&quot;</span> <span class="o">+</span> <span class="s2">&quot;,&quot;</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="nb">str</span><span class="p">(</span><span class="n">i</span><span class="p">)</span> <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">where</span><span class="p">(</span><span class="n">link_matrix_upper</span><span class="p">[</span><span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">,</span> <span class="mi">1</span><span class="p">:]</span> <span class="o">!=</span> <span class="s2">&quot;&quot;</span><span class="p">)[</span><span class="mi">0</span><span class="p">]</span> <span class="o">+</span> <span class="mi">1</span><span class="p">)</span><span class="o">.</span><span class="n">tolist</span><span class="p">())</span>
                <span class="c1"># Lags upto tau_max</span>
                <span class="c1">#autodep_lags = np.argsort(val_matrix[u, v][1:])[::-1]</span>
                <span class="c1">#autodep_lags += 1</span>
                <span class="c1">#autodeplags[u] = &quot;\n\n\n&quot; + &quot;,&quot;.join(str(i) for i in autodep_lags.tolist())</span>

            <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;inner_edge_attribute&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span>
            <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;outer_edge_attribute&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span>

        <span class="c1"># dic[&#39;outer_edge_edge&#39;] = False</span>
        <span class="c1"># dic[&#39;outer_edge_edgecolor&#39;] = None</span>
        <span class="c1"># dic[&#39;inner_edge_edge&#39;] = False</span>
        <span class="c1"># dic[&#39;inner_edge_edgecolor&#39;] = None</span>

    <span class="k">if</span> <span class="n">special_nodes</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">special_nodes_draw</span> <span class="o">=</span> <span class="p">{}</span>
        <span class="k">for</span> <span class="n">node</span> <span class="ow">in</span> <span class="n">special_nodes</span><span class="p">:</span>
            <span class="n">i</span><span class="p">,</span> <span class="n">tau</span> <span class="o">=</span> <span class="n">node</span>
            <span class="k">if</span> <span class="n">tau</span> <span class="o">&gt;=</span> <span class="o">-</span><span class="n">tau_max</span><span class="p">:</span>
                <span class="n">special_nodes_draw</span><span class="p">[</span><span class="n">i</span><span class="p">]</span> <span class="o">=</span> <span class="n">special_nodes</span><span class="p">[</span><span class="n">node</span><span class="p">]</span>
        <span class="n">special_nodes</span> <span class="o">=</span> <span class="n">special_nodes_draw</span>
    

    <span class="c1"># If no links are present, set value to zero</span>
    <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">all_strengths</span><span class="p">)</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
        <span class="n">all_strengths</span> <span class="o">=</span> <span class="p">[</span><span class="mf">0.0</span><span class="p">]</span>

    <span class="k">if</span> <span class="n">node_pos</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">pos</span> <span class="o">=</span> <span class="n">nx</span><span class="o">.</span><span class="n">circular_layout</span><span class="p">(</span><span class="n">deepcopy</span><span class="p">(</span><span class="n">G</span><span class="p">))</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">pos</span> <span class="o">=</span> <span class="p">{}</span>
        <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">N</span><span class="p">):</span>
            <span class="n">pos</span><span class="p">[</span><span class="n">i</span><span class="p">]</span> <span class="o">=</span> <span class="p">(</span><span class="n">node_pos</span><span class="p">[</span><span class="s2">&quot;x&quot;</span><span class="p">][</span><span class="n">i</span><span class="p">],</span> <span class="n">node_pos</span><span class="p">[</span><span class="s2">&quot;y&quot;</span><span class="p">][</span><span class="n">i</span><span class="p">])</span>

    <span class="k">if</span> <span class="n">node_pos</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> <span class="s1">&#39;transform&#39;</span> <span class="ow">in</span> <span class="n">node_pos</span><span class="p">:</span> 
        <span class="n">transform</span> <span class="o">=</span> <span class="n">node_pos</span><span class="p">[</span><span class="s1">&#39;transform&#39;</span><span class="p">]</span>
    <span class="k">else</span><span class="p">:</span> <span class="n">transform</span> <span class="o">=</span> <span class="n">ax</span><span class="o">.</span><span class="n">transData</span>

    <span class="k">if</span> <span class="n">cmap_nodes</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">node_color</span> <span class="o">=</span> <span class="kc">None</span>

    <span class="n">node_rings</span> <span class="o">=</span> <span class="p">{</span>
        <span class="mi">0</span><span class="p">:</span> <span class="p">{</span>
            <span class="s2">&quot;sizes&quot;</span><span class="p">:</span> <span class="kc">None</span><span class="p">,</span>
            <span class="s2">&quot;color_array&quot;</span><span class="p">:</span> <span class="n">node_color</span><span class="p">,</span>
            <span class="s2">&quot;cmap&quot;</span><span class="p">:</span> <span class="n">cmap_nodes</span><span class="p">,</span>
            <span class="s2">&quot;vmin&quot;</span><span class="p">:</span> <span class="n">vmin_nodes</span><span class="p">,</span>
            <span class="s2">&quot;vmax&quot;</span><span class="p">:</span> <span class="n">vmax_nodes</span><span class="p">,</span>
            <span class="s2">&quot;ticks&quot;</span><span class="p">:</span> <span class="n">node_ticks</span><span class="p">,</span>
            <span class="s2">&quot;label&quot;</span><span class="p">:</span> <span class="n">node_colorbar_label</span><span class="p">,</span>
            <span class="s2">&quot;colorbar&quot;</span><span class="p">:</span> <span class="n">show_colorbar</span><span class="p">,</span>
        <span class="p">}</span>
    <span class="p">}</span>

    <span class="n">_draw_network_with_curved_edges</span><span class="p">(</span>
        <span class="n">fig</span><span class="o">=</span><span class="n">fig</span><span class="p">,</span>
        <span class="n">ax</span><span class="o">=</span><span class="n">ax</span><span class="p">,</span>
        <span class="n">G</span><span class="o">=</span><span class="n">deepcopy</span><span class="p">(</span><span class="n">G</span><span class="p">),</span>
        <span class="n">pos</span><span class="o">=</span><span class="n">pos</span><span class="p">,</span>
        <span class="c1"># dictionary of rings: {0:{&#39;sizes&#39;:(N,)-array, &#39;color_array&#39;:(N,)-array</span>
        <span class="c1"># or None, &#39;cmap&#39;:string,</span>
        <span class="n">node_rings</span><span class="o">=</span><span class="n">node_rings</span><span class="p">,</span>
        <span class="c1"># &#39;vmin&#39;:float or None, &#39;vmax&#39;:float or None, &#39;label&#39;:string or None}}</span>
        <span class="n">node_labels</span><span class="o">=</span><span class="n">var_names</span><span class="p">,</span>
        <span class="n">node_label_size</span><span class="o">=</span><span class="n">node_label_size</span><span class="p">,</span>
        <span class="n">node_alpha</span><span class="o">=</span><span class="n">alpha</span><span class="p">,</span>
        <span class="n">standard_size</span><span class="o">=</span><span class="n">node_size</span><span class="p">,</span>
        <span class="n">node_aspect</span><span class="o">=</span><span class="n">node_aspect</span><span class="p">,</span>
        <span class="n">standard_cmap</span><span class="o">=</span><span class="s2">&quot;OrRd&quot;</span><span class="p">,</span>
        <span class="n">standard_color_nodes</span><span class="o">=</span><span class="s2">&quot;lightgrey&quot;</span><span class="p">,</span>
        <span class="n">standard_color_links</span><span class="o">=</span><span class="s2">&quot;black&quot;</span><span class="p">,</span>
        <span class="n">log_sizes</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
        <span class="n">cmap_links</span><span class="o">=</span><span class="n">cmap_edges</span><span class="p">,</span>
        <span class="n">links_vmin</span><span class="o">=</span><span class="n">vmin_edges</span><span class="p">,</span>
        <span class="n">links_vmax</span><span class="o">=</span><span class="n">vmax_edges</span><span class="p">,</span>
        <span class="n">links_ticks</span><span class="o">=</span><span class="n">edge_ticks</span><span class="p">,</span>
        <span class="n">tick_label_size</span><span class="o">=</span><span class="n">tick_label_size</span><span class="p">,</span>
        <span class="c1"># cmap_links_edges=&#39;YlOrRd&#39;, links_edges_vmin=-1., links_edges_vmax=1.,</span>
        <span class="c1"># links_edges_ticks=.2, link_edge_colorbar_label=&#39;link_edge&#39;,</span>
        <span class="n">arrowstyle</span><span class="o">=</span><span class="s2">&quot;simple&quot;</span><span class="p">,</span>
        <span class="n">arrowhead_size</span><span class="o">=</span><span class="n">arrowhead_size</span><span class="p">,</span>
        <span class="n">curved_radius</span><span class="o">=</span><span class="n">curved_radius</span><span class="p">,</span>
        <span class="n">label_fontsize</span><span class="o">=</span><span class="n">label_fontsize</span><span class="p">,</span>
        <span class="n">link_label_fontsize</span><span class="o">=</span><span class="n">link_label_fontsize</span><span class="p">,</span>
        <span class="n">link_colorbar_label</span><span class="o">=</span><span class="n">link_colorbar_label</span><span class="p">,</span>
        <span class="c1"># network_lower_bound=network_lower_bound,</span>
        <span class="n">show_colorbar</span><span class="o">=</span><span class="n">show_colorbar</span><span class="p">,</span>
        <span class="c1"># label_fraction=label_fraction,</span>
        <span class="n">special_nodes</span><span class="o">=</span><span class="n">special_nodes</span><span class="p">,</span>
        <span class="n">autodep_sig_lags</span><span class="o">=</span><span class="n">autodep_sig_lags</span><span class="p">,</span>
        <span class="n">show_autodependency_lags</span><span class="o">=</span><span class="n">show_autodependency_lags</span><span class="p">,</span>
        <span class="n">transform</span><span class="o">=</span><span class="n">transform</span>
    <span class="p">)</span>

    <span class="k">if</span> <span class="n">save_name</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">pyplot</span><span class="o">.</span><span class="n">savefig</span><span class="p">(</span><span class="n">save_name</span><span class="p">,</span> <span class="n">dpi</span><span class="o">=</span><span class="mi">300</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">fig</span><span class="p">,</span> <span class="n">ax</span></div>


<span class="k">def</span> <span class="nf">_reverse_patt</span><span class="p">(</span><span class="n">patt</span><span class="p">):</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;Inverts a link pattern&quot;&quot;&quot;</span>

    <span class="k">if</span> <span class="n">patt</span> <span class="o">==</span> <span class="s2">&quot;&quot;</span><span class="p">:</span>
        <span class="k">return</span> <span class="s2">&quot;&quot;</span>

    <span class="n">left_mark</span><span class="p">,</span> <span class="n">middle_mark</span><span class="p">,</span> <span class="n">right_mark</span> <span class="o">=</span> <span class="n">patt</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">patt</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="n">patt</span><span class="p">[</span><span class="mi">2</span><span class="p">]</span>
    <span class="k">if</span> <span class="n">left_mark</span> <span class="o">==</span> <span class="s2">&quot;&lt;&quot;</span><span class="p">:</span>
        <span class="n">new_right_mark</span> <span class="o">=</span> <span class="s2">&quot;&gt;&quot;</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">new_right_mark</span> <span class="o">=</span> <span class="n">left_mark</span>
    <span class="k">if</span> <span class="n">right_mark</span> <span class="o">==</span> <span class="s2">&quot;&gt;&quot;</span><span class="p">:</span>
        <span class="n">new_left_mark</span> <span class="o">=</span> <span class="s2">&quot;&lt;&quot;</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">new_left_mark</span> <span class="o">=</span> <span class="n">right_mark</span>

    <span class="k">return</span> <span class="n">new_left_mark</span> <span class="o">+</span> <span class="n">middle_mark</span> <span class="o">+</span> <span class="n">new_right_mark</span>

    <span class="c1"># if patt in [&#39;---&#39;, &#39;o--&#39;, &#39;--o&#39;, &#39;o-o&#39;, &#39;&#39;]:</span>
    <span class="c1">#     return patt[::-1]</span>
    <span class="c1"># elif patt == &#39;&lt;-&gt;&#39;:</span>
    <span class="c1">#     return &#39;&lt;-&gt;&#39;</span>
    <span class="c1"># elif patt == &#39;o-&gt;&#39;:</span>
    <span class="c1">#     return &#39;&lt;-o&#39;</span>
    <span class="c1"># elif patt == &#39;&lt;-o&#39;:</span>
    <span class="c1">#     return &#39;o-&gt;&#39;</span>
    <span class="c1"># elif patt == &#39;--&gt;&#39;:</span>
    <span class="c1">#     return &#39;&lt;--&#39;</span>
    <span class="c1"># elif patt == &#39;&lt;--&#39;:</span>
    <span class="c1">#     return &#39;--&gt;&#39;</span>


<span class="k">def</span> <span class="nf">_check_matrices</span><span class="p">(</span><span class="n">graph</span><span class="p">,</span> <span class="n">val_matrix</span><span class="p">,</span> <span class="n">link_width</span><span class="p">,</span> <span class="n">link_attribute</span><span class="p">):</span>

    <span class="k">if</span> <span class="n">graph</span><span class="o">.</span><span class="n">dtype</span> <span class="o">!=</span> <span class="s2">&quot;&lt;U3&quot;</span><span class="p">:</span>
        <span class="c1"># Transform to new graph data type U3</span>
        <span class="n">old_matrix</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">copy</span><span class="p">(</span><span class="n">graph</span><span class="p">)</span>
        <span class="n">graph</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="n">old_matrix</span><span class="o">.</span><span class="n">shape</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="s2">&quot;&lt;U3&quot;</span><span class="p">)</span>
        <span class="n">graph</span><span class="p">[:]</span> <span class="o">=</span> <span class="s2">&quot;&quot;</span>
        <span class="k">for</span> <span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">,</span> <span class="n">tau</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="o">*</span><span class="n">np</span><span class="o">.</span><span class="n">where</span><span class="p">(</span><span class="n">old_matrix</span><span class="p">)):</span>
            <span class="k">if</span> <span class="n">tau</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
                <span class="k">if</span> <span class="n">old_matrix</span><span class="p">[</span><span class="n">j</span><span class="p">,</span> <span class="n">i</span><span class="p">,</span> <span class="mi">0</span><span class="p">]</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
                    <span class="n">graph</span><span class="p">[</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">,</span> <span class="mi">0</span><span class="p">]</span> <span class="o">=</span> <span class="s2">&quot;--&gt;&quot;</span>
                    <span class="n">graph</span><span class="p">[</span><span class="n">j</span><span class="p">,</span> <span class="n">i</span><span class="p">,</span> <span class="mi">0</span><span class="p">]</span> <span class="o">=</span> <span class="s2">&quot;&lt;--&quot;</span>
                <span class="k">else</span><span class="p">:</span>
                    <span class="n">graph</span><span class="p">[</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">,</span> <span class="mi">0</span><span class="p">]</span> <span class="o">=</span> <span class="s2">&quot;o-o&quot;</span>
                    <span class="n">graph</span><span class="p">[</span><span class="n">j</span><span class="p">,</span> <span class="n">i</span><span class="p">,</span> <span class="mi">0</span><span class="p">]</span> <span class="o">=</span> <span class="s2">&quot;o-o&quot;</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">graph</span><span class="p">[</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">,</span> <span class="n">tau</span><span class="p">]</span> <span class="o">=</span> <span class="s2">&quot;--&gt;&quot;</span>
    <span class="k">if</span> <span class="n">graph</span><span class="o">.</span><span class="n">ndim</span> <span class="o">==</span> <span class="mi">4</span><span class="p">:</span>
        <span class="k">for</span> <span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">,</span> <span class="n">taui</span><span class="p">,</span> <span class="n">tauj</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="o">*</span><span class="n">np</span><span class="o">.</span><span class="n">where</span><span class="p">(</span><span class="n">graph</span><span class="p">)):</span>
            <span class="k">if</span> <span class="n">graph</span><span class="p">[</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">,</span> <span class="n">taui</span><span class="p">,</span> <span class="n">tauj</span><span class="p">]</span> <span class="ow">not</span> <span class="ow">in</span> <span class="p">[</span>
                <span class="s2">&quot;---&quot;</span><span class="p">,</span>
                <span class="s2">&quot;o--&quot;</span><span class="p">,</span>
                <span class="s2">&quot;--o&quot;</span><span class="p">,</span>
                <span class="s2">&quot;o-o&quot;</span><span class="p">,</span>
                <span class="s2">&quot;o-&gt;&quot;</span><span class="p">,</span>
                <span class="s2">&quot;&lt;-o&quot;</span><span class="p">,</span>
                <span class="s2">&quot;--&gt;&quot;</span><span class="p">,</span>
                <span class="s2">&quot;&lt;--&quot;</span><span class="p">,</span>
                <span class="s2">&quot;&lt;-&gt;&quot;</span><span class="p">,</span>
                <span class="s2">&quot;x-o&quot;</span><span class="p">,</span>
                <span class="s2">&quot;o-x&quot;</span><span class="p">,</span>
                <span class="s2">&quot;x--&quot;</span><span class="p">,</span>
                <span class="s2">&quot;--x&quot;</span><span class="p">,</span>
                <span class="s2">&quot;x-&gt;&quot;</span><span class="p">,</span>
                <span class="s2">&quot;&lt;-x&quot;</span><span class="p">,</span>
                <span class="s2">&quot;x-x&quot;</span><span class="p">,</span>
                <span class="s2">&quot;&lt;-+&quot;</span><span class="p">,</span>
                <span class="s2">&quot;+-&gt;&quot;</span><span class="p">,</span>
            <span class="p">]:</span>
                <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;Invalid graph entry.&quot;</span><span class="p">)</span>
            <span class="k">if</span> <span class="n">graph</span><span class="p">[</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">,</span> <span class="n">taui</span><span class="p">,</span> <span class="n">tauj</span><span class="p">]</span> <span class="o">!=</span> <span class="n">_reverse_patt</span><span class="p">(</span><span class="n">graph</span><span class="p">[</span><span class="n">j</span><span class="p">,</span> <span class="n">i</span><span class="p">,</span> <span class="n">tauj</span><span class="p">,</span> <span class="n">taui</span><span class="p">]):</span>
                <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span>
                    <span class="s2">&quot;graph needs to have consistent entries: &quot;</span>
                    <span class="s2">&quot;graph[i, j, taui, tauj] == _reverse_patt(graph[j, i, tauj, taui])&quot;</span><span class="p">)</span>
            <span class="k">if</span> <span class="p">(</span>
                <span class="n">val_matrix</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span>
                <span class="ow">and</span> <span class="n">val_matrix</span><span class="p">[</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">,</span> <span class="n">taui</span><span class="p">,</span> <span class="n">tauj</span><span class="p">]</span> <span class="o">!=</span> <span class="n">val_matrix</span><span class="p">[</span><span class="n">j</span><span class="p">,</span> <span class="n">i</span><span class="p">,</span> <span class="n">tauj</span><span class="p">,</span> <span class="n">taui</span><span class="p">]</span>
            <span class="p">):</span>
                <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span>
                    <span class="s2">&quot;val_matrix needs to have consistent entries: &quot;</span>
                    <span class="s2">&quot;val_matrix[i, j, taui, tauj] == val_matrix[j, i, tauj, taui]&quot;</span><span class="p">)</span>
            <span class="k">if</span> <span class="p">(</span>
                <span class="n">link_width</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span>
                <span class="ow">and</span> <span class="n">link_width</span><span class="p">[</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">,</span> <span class="n">taui</span><span class="p">,</span> <span class="n">tauj</span><span class="p">]</span> <span class="o">!=</span> <span class="n">link_width</span><span class="p">[</span><span class="n">j</span><span class="p">,</span> <span class="n">i</span><span class="p">,</span> <span class="n">tauj</span><span class="p">,</span> <span class="n">taui</span><span class="p">]</span>
            <span class="p">):</span>
                <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span>
                    <span class="s2">&quot;link_width needs to have consistent entries: &quot;</span>
                    <span class="s2">&quot;link_width[i, j, taui, tauj] == link_width[j, i, tauj, taui]&quot;</span><span class="p">)</span>            
            <span class="k">if</span> <span class="p">(</span>
                <span class="n">link_attribute</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span>
                <span class="ow">and</span> <span class="n">link_attribute</span><span class="p">[</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">,</span> <span class="n">taui</span><span class="p">,</span> <span class="n">tauj</span><span class="p">]</span> <span class="o">!=</span> <span class="n">link_attribute</span><span class="p">[</span><span class="n">j</span><span class="p">,</span> <span class="n">i</span><span class="p">,</span> <span class="n">tauj</span><span class="p">,</span> <span class="n">taui</span><span class="p">]</span>
            <span class="p">):</span>
                <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span>
                    <span class="s2">&quot;link_attribute needs to have consistent entries: &quot;</span>
                    <span class="s2">&quot;link_attribute[i, j, taui, tauj] == link_attribute[j, i, tauj, taui]&quot;</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="c1"># print(graph[:,:,0])</span>
        <span class="c1"># Assert that graph has valid and consistent lag-zero entries</span>
        <span class="k">for</span> <span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">,</span> <span class="n">tau</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="o">*</span><span class="n">np</span><span class="o">.</span><span class="n">where</span><span class="p">(</span><span class="n">graph</span><span class="p">)):</span>
            <span class="k">if</span> <span class="n">tau</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
                <span class="k">if</span> <span class="n">graph</span><span class="p">[</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">,</span> <span class="mi">0</span><span class="p">]</span> <span class="o">!=</span> <span class="n">_reverse_patt</span><span class="p">(</span><span class="n">graph</span><span class="p">[</span><span class="n">j</span><span class="p">,</span> <span class="n">i</span><span class="p">,</span> <span class="mi">0</span><span class="p">]):</span>
                    <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span>
                        <span class="s2">&quot;graph needs to have consistent lag-zero links, but &quot;</span>
                        <span class="s2">&quot; graph[</span><span class="si">%d</span><span class="s2">,</span><span class="si">%d</span><span class="s2">,0]=</span><span class="si">%s</span><span class="s2"> and graph[</span><span class="si">%d</span><span class="s2">,</span><span class="si">%d</span><span class="s2">,0]=</span><span class="si">%s</span><span class="s2">)&quot;</span> <span class="o">%</span><span class="p">(</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">,</span> <span class="n">graph</span><span class="p">[</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">,</span> <span class="mi">0</span><span class="p">],</span> <span class="n">j</span><span class="p">,</span> <span class="n">i</span><span class="p">,</span> <span class="n">graph</span><span class="p">[</span><span class="n">j</span><span class="p">,</span> <span class="n">i</span><span class="p">,</span> <span class="mi">0</span><span class="p">])</span>
                    <span class="p">)</span>
                <span class="k">if</span> <span class="p">(</span>
                    <span class="n">val_matrix</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span>
                    <span class="ow">and</span> <span class="n">val_matrix</span><span class="p">[</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">,</span> <span class="mi">0</span><span class="p">]</span> <span class="o">!=</span> <span class="n">val_matrix</span><span class="p">[</span><span class="n">j</span><span class="p">,</span> <span class="n">i</span><span class="p">,</span> <span class="mi">0</span><span class="p">]</span>
                <span class="p">):</span>
                    <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;val_matrix needs to be symmetric for lag-zero&quot;</span><span class="p">)</span>
                <span class="k">if</span> <span class="p">(</span>
                    <span class="n">link_width</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span>
                    <span class="ow">and</span> <span class="n">link_width</span><span class="p">[</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">,</span> <span class="mi">0</span><span class="p">]</span> <span class="o">!=</span> <span class="n">link_width</span><span class="p">[</span><span class="n">j</span><span class="p">,</span> <span class="n">i</span><span class="p">,</span> <span class="mi">0</span><span class="p">]</span>
                <span class="p">):</span>
                    <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;link_width needs to be symmetric for lag-zero&quot;</span><span class="p">)</span>
                <span class="k">if</span> <span class="p">(</span>
                    <span class="n">link_attribute</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span>
                    <span class="ow">and</span> <span class="n">link_attribute</span><span class="p">[</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">,</span> <span class="mi">0</span><span class="p">]</span> <span class="o">!=</span> <span class="n">link_attribute</span><span class="p">[</span><span class="n">j</span><span class="p">,</span> <span class="n">i</span><span class="p">,</span> <span class="mi">0</span><span class="p">]</span>
                <span class="p">):</span>
                    <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span>
                        <span class="s2">&quot;link_attribute needs to be symmetric for lag-zero&quot;</span>
                    <span class="p">)</span>

            <span class="k">if</span> <span class="n">graph</span><span class="p">[</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">,</span> <span class="n">tau</span><span class="p">]</span> <span class="ow">not</span> <span class="ow">in</span> <span class="p">[</span>
                <span class="s2">&quot;---&quot;</span><span class="p">,</span>
                <span class="s2">&quot;o--&quot;</span><span class="p">,</span>
                <span class="s2">&quot;--o&quot;</span><span class="p">,</span>
                <span class="s2">&quot;o-o&quot;</span><span class="p">,</span>
                <span class="s2">&quot;o-&gt;&quot;</span><span class="p">,</span>
                <span class="s2">&quot;&lt;-o&quot;</span><span class="p">,</span>
                <span class="s2">&quot;--&gt;&quot;</span><span class="p">,</span>
                <span class="s2">&quot;&lt;--&quot;</span><span class="p">,</span>
                <span class="s2">&quot;&lt;-&gt;&quot;</span><span class="p">,</span>
                <span class="s2">&quot;x-o&quot;</span><span class="p">,</span>
                <span class="s2">&quot;o-x&quot;</span><span class="p">,</span>
                <span class="s2">&quot;x--&quot;</span><span class="p">,</span>
                <span class="s2">&quot;--x&quot;</span><span class="p">,</span>
                <span class="s2">&quot;x-&gt;&quot;</span><span class="p">,</span>
                <span class="s2">&quot;&lt;-x&quot;</span><span class="p">,</span>
                <span class="s2">&quot;x-x&quot;</span><span class="p">,</span>
                <span class="s2">&quot;&lt;-+&quot;</span><span class="p">,</span>
                <span class="s2">&quot;+-&gt;&quot;</span><span class="p">,</span>
            <span class="p">]:</span>
                <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;Invalid graph entry.&quot;</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">val_matrix</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="c1"># if graph.ndim == 4:</span>
        <span class="c1">#     val_matrix = (graph != &quot;&quot;).astype(&quot;int&quot;)</span>
        <span class="c1"># else:</span>
            <span class="n">val_matrix</span> <span class="o">=</span> <span class="p">(</span><span class="n">graph</span> <span class="o">!=</span> <span class="s2">&quot;&quot;</span><span class="p">)</span><span class="o">.</span><span class="n">astype</span><span class="p">(</span><span class="s2">&quot;int&quot;</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">link_width</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> <span class="ow">not</span> <span class="n">np</span><span class="o">.</span><span class="n">all</span><span class="p">(</span><span class="n">link_width</span> <span class="o">&gt;=</span> <span class="mf">0.0</span><span class="p">):</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;link_width must be non-negative&quot;</span><span class="p">)</span>

    <span class="k">return</span> <span class="n">graph</span><span class="p">,</span> <span class="n">val_matrix</span><span class="p">,</span> <span class="n">link_width</span><span class="p">,</span> <span class="n">link_attribute</span>


<div class="viewcode-block" id="plot_time_series_graph"><a class="viewcode-back" href="../../index.html#tigramite.plotting.plot_time_series_graph">[docs]</a><span class="k">def</span> <span class="nf">plot_time_series_graph</span><span class="p">(</span>
        <span class="n">graph</span><span class="p">,</span>
        <span class="n">val_matrix</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
        <span class="n">var_names</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
        <span class="n">fig_ax</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
        <span class="n">figsize</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
        <span class="n">link_colorbar_label</span><span class="o">=</span><span class="s2">&quot;MCI&quot;</span><span class="p">,</span>
        <span class="n">save_name</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
        <span class="n">link_width</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
        <span class="n">link_attribute</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
        <span class="n">arrow_linewidth</span><span class="o">=</span><span class="mi">4</span><span class="p">,</span>
        <span class="n">vmin_edges</span><span class="o">=-</span><span class="mi">1</span><span class="p">,</span>
        <span class="n">vmax_edges</span><span class="o">=</span><span class="mf">1.0</span><span class="p">,</span>
        <span class="n">edge_ticks</span><span class="o">=</span><span class="mf">0.4</span><span class="p">,</span>
        <span class="n">cmap_edges</span><span class="o">=</span><span class="s2">&quot;RdBu_r&quot;</span><span class="p">,</span>
        <span class="n">order</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
        <span class="n">node_size</span><span class="o">=</span><span class="mf">0.1</span><span class="p">,</span>
        <span class="n">node_aspect</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
        <span class="n">arrowhead_size</span><span class="o">=</span><span class="mi">20</span><span class="p">,</span>
        <span class="n">curved_radius</span><span class="o">=</span><span class="mf">0.2</span><span class="p">,</span>
        <span class="n">label_fontsize</span><span class="o">=</span><span class="mi">10</span><span class="p">,</span>
        <span class="n">tick_label_size</span><span class="o">=</span><span class="mi">6</span><span class="p">,</span>
        <span class="n">alpha</span><span class="o">=</span><span class="mf">1.0</span><span class="p">,</span>
        <span class="n">inner_edge_style</span><span class="o">=</span><span class="s2">&quot;dashed&quot;</span><span class="p">,</span>
        <span class="n">link_matrix</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
        <span class="n">special_nodes</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
        <span class="n">node_classification</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
        <span class="c1"># aux_graph=None,</span>
        <span class="n">standard_color_links</span><span class="o">=</span><span class="s1">&#39;black&#39;</span><span class="p">,</span>
        <span class="n">standard_color_nodes</span><span class="o">=</span><span class="s1">&#39;lightgrey&#39;</span><span class="p">,</span>
<span class="p">):</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;Creates a time series graph.</span>
<span class="sd">    This is still in beta. The time series graph&#39;s links are colored by</span>
<span class="sd">    val_matrix.</span>

<span class="sd">    Parameters</span>
<span class="sd">    ----------</span>
<span class="sd">    graph : string or bool array-like, optional (default: None)</span>
<span class="sd">        Either string matrix providing graph or bool array providing only adjacencies</span>
<span class="sd">        Either of shape (N, N, tau_max + 1) or as auxiliary graph of dims </span>
<span class="sd">        (N, N, tau_max+1, tau_max+1) describing auxADMG. </span>
<span class="sd">    val_matrix : array_like</span>
<span class="sd">        Matrix of same shape as graph containing test statistic values.</span>
<span class="sd">    var_names : list, optional (default: None)</span>
<span class="sd">        List of variable names. If None, range(N) is used.</span>
<span class="sd">    fig_ax : tuple of figure and axis object, optional (default: None)</span>
<span class="sd">        Figure and axes instance. If None they are created.</span>
<span class="sd">    figsize : tuple</span>
<span class="sd">        Size of figure.</span>
<span class="sd">    save_name : str, optional (default: None)</span>
<span class="sd">        Name of figure file to save figure. If None, figure is shown in window.</span>
<span class="sd">    link_colorbar_label : str, optional (default: &#39;MCI&#39;)</span>
<span class="sd">        Test statistic label.</span>
<span class="sd">    link_width : array-like, optional (default: None)</span>
<span class="sd">        Array of val_matrix.shape specifying relative link width with maximum</span>
<span class="sd">        given by arrow_linewidth. If None, all links have same width.</span>
<span class="sd">    link_attribute : array-like, optional (default: None)</span>
<span class="sd">        Array of graph.shape specifying specific in drawing the graph (for internal use).</span>
<span class="sd">    order : list, optional (default: None)</span>
<span class="sd">        order of variables from top to bottom.</span>
<span class="sd">    arrow_linewidth : float, optional (default: 30)</span>
<span class="sd">        Linewidth.</span>
<span class="sd">    vmin_edges : float, optional (default: -1)</span>
<span class="sd">        Link colorbar scale lower bound.</span>
<span class="sd">    vmax_edges : float, optional (default: 1)</span>
<span class="sd">        Link colorbar scale upper bound.</span>
<span class="sd">    edge_ticks : float, optional (default: 0.4)</span>
<span class="sd">        Link tick mark interval.</span>
<span class="sd">    cmap_edges : str, optional (default: &#39;RdBu_r&#39;)</span>
<span class="sd">        Colormap for links.</span>
<span class="sd">    node_size : int, optional (default: 0.1)</span>
<span class="sd">        Node size.</span>
<span class="sd">    node_aspect : float, optional (default: None)</span>
<span class="sd">        Ratio between the heigth and width of the varible nodes.</span>
<span class="sd">    arrowhead_size : int, optional (default: 20)</span>
<span class="sd">        Size of link arrow head. Passed on to FancyArrowPatch object.</span>
<span class="sd">    curved_radius, float, optional (default: 0.2)</span>
<span class="sd">        Curvature of links. Passed on to FancyArrowPatch object.</span>
<span class="sd">    label_fontsize : int, optional (default: 10)</span>
<span class="sd">        Fontsize of colorbar labels.</span>
<span class="sd">    alpha : float, optional (default: 1.)</span>
<span class="sd">        Opacity.</span>
<span class="sd">    node_label_size : int, optional (default: 10)</span>
<span class="sd">        Fontsize of node labels.</span>
<span class="sd">    link_label_fontsize : int, optional (default: 6)</span>
<span class="sd">        Fontsize of link labels.</span>
<span class="sd">    tick_label_size : int, optional (default: 6)</span>
<span class="sd">        Fontsize of tick labels.</span>
<span class="sd">    inner_edge_style : string, optional (default: &#39;dashed&#39;)</span>
<span class="sd">        Style of inner_edge contemporaneous links.</span>
<span class="sd">    special_nodes : dict</span>
<span class="sd">        Dictionary of format {(i, -tau): &#39;blue&#39;, ...} to color special nodes.</span>
<span class="sd">    node_classification : dict or None (default: None)</span>
<span class="sd">        Dictionary of format {i: &#39;space_context&#39;, ...} to classify nodes into system, context, or dummy nodes.</span>
<span class="sd">        Keys of the dictionary are from {0, ..., N-1} where N is the number of nodes.</span>
<span class="sd">        Options for the values are &quot;system&quot;, &quot;time_context&quot;, &quot;space_context&quot;, &quot;time_dummy&quot;, or &quot;space_dummy&quot;.</span>
<span class="sd">        Space_contexts and dummy nodes need to be represented as a single node in the time series graph.</span>
<span class="sd">        In case no value is supplied all nodes are treated as system nodes, i.e. are plotted in a time-resolved manner.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">if</span> <span class="n">link_matrix</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;link_matrix is deprecated and replaced by graph array&quot;</span>
                         <span class="s2">&quot; which is now returned by all methods.&quot;</span><span class="p">)</span>
        
    <span class="k">if</span> <span class="n">fig_ax</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">fig</span> <span class="o">=</span> <span class="n">pyplot</span><span class="o">.</span><span class="n">figure</span><span class="p">(</span><span class="n">figsize</span><span class="o">=</span><span class="n">figsize</span><span class="p">)</span>
        <span class="n">ax</span> <span class="o">=</span> <span class="n">fig</span><span class="o">.</span><span class="n">add_subplot</span><span class="p">(</span><span class="mi">111</span><span class="p">,</span> <span class="n">frame_on</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">fig</span><span class="p">,</span> <span class="n">ax</span> <span class="o">=</span> <span class="n">fig_ax</span>

    <span class="k">if</span> <span class="n">val_matrix</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">no_coloring</span> <span class="o">=</span> <span class="kc">True</span>
        <span class="n">cmap_edges</span> <span class="o">=</span> <span class="kc">None</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">no_coloring</span> <span class="o">=</span> <span class="kc">False</span>

    <span class="p">(</span><span class="n">graph</span><span class="p">,</span> <span class="n">val_matrix</span><span class="p">,</span> <span class="n">link_width</span><span class="p">,</span> <span class="n">link_attribute</span><span class="p">)</span> <span class="o">=</span> <span class="n">_check_matrices</span><span class="p">(</span>
        <span class="n">graph</span><span class="p">,</span> <span class="n">val_matrix</span><span class="p">,</span> <span class="n">link_width</span><span class="p">,</span> <span class="n">link_attribute</span>
    <span class="p">)</span>

    <span class="k">if</span> <span class="n">graph</span><span class="o">.</span><span class="n">ndim</span> <span class="o">==</span> <span class="mi">4</span><span class="p">:</span>
        <span class="n">N</span><span class="p">,</span> <span class="n">N</span><span class="p">,</span> <span class="n">dummy</span><span class="p">,</span> <span class="n">_</span> <span class="o">=</span> <span class="n">graph</span><span class="o">.</span><span class="n">shape</span>
        <span class="n">tau_max</span> <span class="o">=</span> <span class="n">dummy</span> <span class="o">-</span> <span class="mi">1</span>
        <span class="n">max_lag</span> <span class="o">=</span> <span class="n">tau_max</span> <span class="o">+</span> <span class="mi">1</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">N</span><span class="p">,</span> <span class="n">N</span><span class="p">,</span> <span class="n">dummy</span> <span class="o">=</span> <span class="n">graph</span><span class="o">.</span><span class="n">shape</span>
        <span class="n">tau_max</span> <span class="o">=</span> <span class="n">dummy</span> <span class="o">-</span> <span class="mi">1</span>
        <span class="n">max_lag</span> <span class="o">=</span> <span class="n">tau_max</span> <span class="o">+</span> <span class="mi">1</span>

    <span class="k">if</span> <span class="n">np</span><span class="o">.</span><span class="n">count_nonzero</span><span class="p">(</span><span class="n">graph</span> <span class="o">==</span> <span class="s2">&quot;&quot;</span><span class="p">)</span> <span class="o">==</span> <span class="n">graph</span><span class="o">.</span><span class="n">size</span><span class="p">:</span>
        <span class="k">if</span> <span class="n">graph</span><span class="o">.</span><span class="n">ndim</span> <span class="o">==</span> <span class="mi">4</span><span class="p">:</span>
            <span class="n">graph</span><span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">]</span> <span class="o">=</span> <span class="s2">&quot;---&quot;</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">graph</span><span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">0</span><span class="p">]</span> <span class="o">=</span> <span class="s2">&quot;---&quot;</span>
        <span class="n">no_links</span> <span class="o">=</span> <span class="kc">True</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">no_links</span> <span class="o">=</span> <span class="kc">False</span>

    <span class="k">if</span> <span class="n">var_names</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">var_names</span> <span class="o">=</span> <span class="nb">range</span><span class="p">(</span><span class="n">N</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">order</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">order</span> <span class="o">=</span> <span class="nb">range</span><span class="p">(</span><span class="n">N</span><span class="p">)</span>

    <span class="k">if</span> <span class="nb">set</span><span class="p">(</span><span class="n">order</span><span class="p">)</span> <span class="o">!=</span> <span class="nb">set</span><span class="p">(</span><span class="nb">range</span><span class="p">(</span><span class="n">N</span><span class="p">)):</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;order must be a permutation of range(N)&quot;</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">translate</span><span class="p">(</span><span class="n">row</span><span class="p">,</span> <span class="n">lag</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">row</span> <span class="o">*</span> <span class="n">max_lag</span> <span class="o">+</span> <span class="n">lag</span>

    <span class="c1"># Define graph links by absolute maximum (positive or negative like for</span>
    <span class="c1"># partial correlation)</span>
    <span class="n">tsg</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="n">N</span> <span class="o">*</span> <span class="n">max_lag</span><span class="p">,</span> <span class="n">N</span> <span class="o">*</span> <span class="n">max_lag</span><span class="p">))</span>
    <span class="n">tsg_val</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="n">N</span> <span class="o">*</span> <span class="n">max_lag</span><span class="p">,</span> <span class="n">N</span> <span class="o">*</span> <span class="n">max_lag</span><span class="p">))</span>
    <span class="n">tsg_width</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="n">N</span> <span class="o">*</span> <span class="n">max_lag</span><span class="p">,</span> <span class="n">N</span> <span class="o">*</span> <span class="n">max_lag</span><span class="p">))</span>
    <span class="n">tsg_style</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="n">N</span> <span class="o">*</span> <span class="n">max_lag</span><span class="p">,</span> <span class="n">N</span> <span class="o">*</span> <span class="n">max_lag</span><span class="p">),</span> <span class="n">dtype</span><span class="o">=</span><span class="n">graph</span><span class="o">.</span><span class="n">dtype</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">link_attribute</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">tsg_attr</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="n">N</span> <span class="o">*</span> <span class="n">max_lag</span><span class="p">,</span> <span class="n">N</span> <span class="o">*</span> <span class="n">max_lag</span><span class="p">),</span> <span class="n">dtype</span><span class="o">=</span><span class="n">link_attribute</span><span class="o">.</span><span class="n">dtype</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">graph</span><span class="o">.</span><span class="n">ndim</span> <span class="o">==</span> <span class="mi">4</span><span class="p">:</span>
        <span class="c1"># 4-dimensional graphs represent the finite-time window projection of stationary 3-d graphs</span>
        <span class="c1"># They are internally created in some classes</span>
        <span class="c1"># Only draw link in one direction</span>
        <span class="k">for</span> <span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">,</span> <span class="n">taui</span><span class="p">,</span> <span class="n">tauj</span> <span class="ow">in</span> <span class="n">np</span><span class="o">.</span><span class="n">column_stack</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">where</span><span class="p">(</span><span class="n">graph</span><span class="p">)):</span>
            <span class="n">tau</span> <span class="o">=</span> <span class="n">taui</span> <span class="o">-</span> <span class="n">tauj</span>
            <span class="c1"># if tau &lt;= 0 and j &lt;= i:</span>
            <span class="k">if</span> <span class="n">translate</span><span class="p">(</span><span class="n">i</span><span class="p">,</span>   <span class="n">max_lag</span> <span class="o">-</span> <span class="mi">1</span> <span class="o">-</span> <span class="n">taui</span><span class="p">)</span> <span class="o">&gt;=</span> <span class="n">translate</span><span class="p">(</span><span class="n">j</span><span class="p">,</span> <span class="n">max_lag</span><span class="o">-</span><span class="mi">1</span><span class="o">-</span><span class="n">tauj</span><span class="p">):</span>
                <span class="k">continue</span>
            <span class="c1"># print(max_lag, (i, -taui), (j, -tauj), aux_graph[i, j, taui, tauj])</span>
            <span class="c1"># print(translate(i, max_lag - 1 - taui), translate(j, max_lag-1-tauj))</span>
            <span class="n">tsg</span><span class="p">[</span><span class="n">translate</span><span class="p">(</span><span class="n">i</span><span class="p">,</span>   <span class="n">max_lag</span> <span class="o">-</span> <span class="mi">1</span> <span class="o">-</span> <span class="n">taui</span><span class="p">),</span> <span class="n">translate</span><span class="p">(</span><span class="n">j</span><span class="p">,</span> <span class="n">max_lag</span><span class="o">-</span><span class="mi">1</span><span class="o">-</span><span class="n">tauj</span><span class="p">)]</span> <span class="o">=</span> <span class="mf">1.0</span>
            <span class="n">tsg_val</span><span class="p">[</span><span class="n">translate</span><span class="p">(</span><span class="n">i</span><span class="p">,</span>   <span class="n">max_lag</span> <span class="o">-</span> <span class="mi">1</span> <span class="o">-</span> <span class="n">taui</span><span class="p">),</span> <span class="n">translate</span><span class="p">(</span><span class="n">j</span><span class="p">,</span> <span class="n">max_lag</span><span class="o">-</span><span class="mi">1</span><span class="o">-</span><span class="n">tauj</span><span class="p">)]</span> <span class="o">=</span> <span class="n">val_matrix</span><span class="p">[</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">,</span> <span class="n">taui</span><span class="p">,</span> <span class="n">tauj</span><span class="p">]</span>
            <span class="n">tsg_style</span><span class="p">[</span><span class="n">translate</span><span class="p">(</span><span class="n">i</span><span class="p">,</span>   <span class="n">max_lag</span> <span class="o">-</span> <span class="mi">1</span> <span class="o">-</span> <span class="n">taui</span><span class="p">),</span> <span class="n">translate</span><span class="p">(</span><span class="n">j</span><span class="p">,</span> <span class="n">max_lag</span><span class="o">-</span><span class="mi">1</span><span class="o">-</span><span class="n">tauj</span><span class="p">)]</span> <span class="o">=</span> <span class="n">graph</span><span class="p">[</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">,</span> <span class="n">taui</span><span class="p">,</span> <span class="n">tauj</span><span class="p">]</span>
            <span class="k">if</span> <span class="n">link_width</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                <span class="n">tsg_width</span><span class="p">[</span><span class="n">translate</span><span class="p">(</span><span class="n">i</span><span class="p">,</span>   <span class="n">max_lag</span> <span class="o">-</span> <span class="mi">1</span> <span class="o">-</span> <span class="n">taui</span><span class="p">),</span> <span class="n">translate</span><span class="p">(</span><span class="n">j</span><span class="p">,</span> <span class="n">max_lag</span><span class="o">-</span><span class="mi">1</span><span class="o">-</span><span class="n">tauj</span><span class="p">)]</span> <span class="o">=</span> <span class="n">link_width</span><span class="p">[</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">,</span> <span class="n">taui</span><span class="p">,</span> <span class="n">tauj</span><span class="p">]</span> <span class="o">/</span> <span class="n">link_width</span><span class="o">.</span><span class="n">max</span><span class="p">()</span> <span class="o">*</span> <span class="n">arrow_linewidth</span>
            <span class="k">if</span> <span class="n">link_attribute</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                <span class="n">tsg_attr</span><span class="p">[</span><span class="n">translate</span><span class="p">(</span><span class="n">i</span><span class="p">,</span>   <span class="n">max_lag</span> <span class="o">-</span> <span class="mi">1</span> <span class="o">-</span> <span class="n">taui</span><span class="p">),</span> <span class="n">translate</span><span class="p">(</span><span class="n">j</span><span class="p">,</span> <span class="n">max_lag</span><span class="o">-</span><span class="mi">1</span><span class="o">-</span><span class="n">tauj</span><span class="p">)]</span> <span class="o">=</span> <span class="n">link_attribute</span><span class="p">[</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">,</span> <span class="n">taui</span><span class="p">,</span> <span class="n">tauj</span><span class="p">]</span> <span class="c1">#&#39;spurious&#39;</span>
        <span class="c1"># print(tsg_style)   </span>
            <span class="c1"># print(tsg_style[translate(i,   max_lag - 1 - taui), translate(j, max_lag-1-tauj)] = graph[i, j, taui, tauj])    </span>
            <span class="c1"># print(max_lag, (i, -taui), (j, -tauj), graph[i, j, taui, tauj], tsg_style[translate(i,   max_lag - 1 - taui), translate(j, max_lag-1-tauj)])</span>
 

    <span class="k">else</span><span class="p">:</span>
      <span class="c1"># Only draw link in one direction</span>
      <span class="c1"># Remove lower triangle</span>
      <span class="n">link_matrix_tsg</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">copy</span><span class="p">(</span><span class="n">graph</span><span class="p">)</span>
      <span class="n">link_matrix_tsg</span><span class="p">[:,</span> <span class="p">:,</span> <span class="mi">0</span><span class="p">]</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">triu</span><span class="p">(</span><span class="n">graph</span><span class="p">[:,</span> <span class="p">:,</span> <span class="mi">0</span><span class="p">])</span>

      <span class="k">for</span> <span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">,</span> <span class="n">tau</span> <span class="ow">in</span> <span class="n">np</span><span class="o">.</span><span class="n">column_stack</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">where</span><span class="p">(</span><span class="n">link_matrix_tsg</span><span class="p">)):</span>
        <span class="k">for</span> <span class="n">t</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">max_lag</span><span class="p">):</span>
            <span class="k">if</span> <span class="p">(</span>
                <span class="mi">0</span> <span class="o">&lt;=</span> <span class="n">translate</span><span class="p">(</span><span class="n">i</span><span class="p">,</span> <span class="n">t</span> <span class="o">-</span> <span class="n">tau</span><span class="p">)</span>
                <span class="ow">and</span> <span class="n">translate</span><span class="p">(</span><span class="n">i</span><span class="p">,</span> <span class="n">t</span> <span class="o">-</span> <span class="n">tau</span><span class="p">)</span> <span class="o">%</span> <span class="n">max_lag</span> <span class="o">&lt;=</span> <span class="n">translate</span><span class="p">(</span><span class="n">j</span><span class="p">,</span> <span class="n">t</span><span class="p">)</span> <span class="o">%</span> <span class="n">max_lag</span>
            <span class="p">):</span>

                <span class="n">tsg</span><span class="p">[</span>
                    <span class="n">translate</span><span class="p">(</span><span class="n">i</span><span class="p">,</span> <span class="n">t</span> <span class="o">-</span> <span class="n">tau</span><span class="p">),</span> <span class="n">translate</span><span class="p">(</span><span class="n">j</span><span class="p">,</span> <span class="n">t</span><span class="p">)</span>
                <span class="p">]</span> <span class="o">=</span> <span class="mf">1.0</span>  <span class="c1"># val_matrix[i, j, tau]</span>
                <span class="n">tsg_val</span><span class="p">[</span><span class="n">translate</span><span class="p">(</span><span class="n">i</span><span class="p">,</span> <span class="n">t</span> <span class="o">-</span> <span class="n">tau</span><span class="p">),</span> <span class="n">translate</span><span class="p">(</span><span class="n">j</span><span class="p">,</span> <span class="n">t</span><span class="p">)]</span> <span class="o">=</span> <span class="n">val_matrix</span><span class="p">[</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">,</span> <span class="n">tau</span><span class="p">]</span>
                <span class="n">tsg_style</span><span class="p">[</span><span class="n">translate</span><span class="p">(</span><span class="n">i</span><span class="p">,</span> <span class="n">t</span> <span class="o">-</span> <span class="n">tau</span><span class="p">),</span> <span class="n">translate</span><span class="p">(</span><span class="n">j</span><span class="p">,</span> <span class="n">t</span><span class="p">)]</span> <span class="o">=</span> <span class="n">graph</span><span class="p">[</span>
                    <span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">,</span> <span class="n">tau</span>
                <span class="p">]</span>
                <span class="k">if</span> <span class="n">link_width</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                    <span class="n">tsg_width</span><span class="p">[</span><span class="n">translate</span><span class="p">(</span><span class="n">i</span><span class="p">,</span> <span class="n">t</span> <span class="o">-</span> <span class="n">tau</span><span class="p">),</span> <span class="n">translate</span><span class="p">(</span><span class="n">j</span><span class="p">,</span> <span class="n">t</span><span class="p">)]</span> <span class="o">=</span> <span class="p">(</span>
                        <span class="n">link_width</span><span class="p">[</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">,</span> <span class="n">tau</span><span class="p">]</span> <span class="o">/</span> <span class="n">link_width</span><span class="o">.</span><span class="n">max</span><span class="p">()</span> <span class="o">*</span> <span class="n">arrow_linewidth</span>
                    <span class="p">)</span>
                <span class="k">if</span> <span class="n">link_attribute</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                    <span class="n">tsg_attr</span><span class="p">[</span><span class="n">translate</span><span class="p">(</span><span class="n">i</span><span class="p">,</span> <span class="n">t</span> <span class="o">-</span> <span class="n">tau</span><span class="p">),</span> <span class="n">translate</span><span class="p">(</span><span class="n">j</span><span class="p">,</span> <span class="n">t</span><span class="p">)]</span> <span class="o">=</span> <span class="n">link_attribute</span><span class="p">[</span>
                        <span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">,</span> <span class="n">tau</span>
                    <span class="p">]</span>


    <span class="n">G</span> <span class="o">=</span> <span class="n">nx</span><span class="o">.</span><span class="n">DiGraph</span><span class="p">(</span><span class="n">tsg</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">special_nodes</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">special_nodes_tsg</span> <span class="o">=</span> <span class="p">{}</span>
        <span class="k">for</span> <span class="n">node</span> <span class="ow">in</span> <span class="n">special_nodes</span><span class="p">:</span>
            <span class="n">i</span><span class="p">,</span> <span class="n">tau</span> <span class="o">=</span> <span class="n">node</span>
            <span class="k">if</span> <span class="n">tau</span> <span class="o">&gt;=</span> <span class="o">-</span><span class="n">tau_max</span><span class="p">:</span>
                <span class="n">special_nodes_tsg</span><span class="p">[</span><span class="n">translate</span><span class="p">(</span><span class="n">i</span><span class="p">,</span> <span class="n">max_lag</span><span class="o">-</span><span class="mi">1</span> <span class="o">+</span> <span class="n">tau</span><span class="p">)]</span> <span class="o">=</span> <span class="n">special_nodes</span><span class="p">[</span><span class="n">node</span><span class="p">]</span>

        <span class="n">special_nodes</span> <span class="o">=</span> <span class="n">special_nodes_tsg</span>

    <span class="k">if</span> <span class="n">node_classification</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">node_classification</span> <span class="o">=</span> <span class="p">{</span><span class="n">i</span><span class="p">:</span> <span class="s2">&quot;system&quot;</span> <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">N</span><span class="p">)}</span>
    <span class="n">node_classification_tsg</span> <span class="o">=</span> <span class="p">{}</span>
    <span class="k">for</span> <span class="n">node</span> <span class="ow">in</span> <span class="n">node_classification</span><span class="p">:</span>
        <span class="k">for</span> <span class="n">tau</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">max_lag</span><span class="p">):</span>
            <span class="k">if</span> <span class="n">tau</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
                <span class="n">suffix</span> <span class="o">=</span> <span class="s2">&quot;_first&quot;</span>
            <span class="k">elif</span> <span class="n">tau</span> <span class="o">==</span> <span class="n">max_lag</span><span class="o">-</span><span class="mi">1</span><span class="p">:</span>
                <span class="n">suffix</span> <span class="o">=</span> <span class="s2">&quot;_last&quot;</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">suffix</span> <span class="o">=</span> <span class="s2">&quot;_middle&quot;</span>
            <span class="n">node_classification_tsg</span><span class="p">[</span><span class="n">translate</span><span class="p">(</span><span class="n">node</span><span class="p">,</span> <span class="n">tau</span><span class="p">)]</span> <span class="o">=</span> <span class="n">node_classification</span><span class="p">[</span><span class="n">node</span><span class="p">]</span> <span class="o">+</span> <span class="n">suffix</span>

    <span class="c1"># node_color = np.zeros(N)</span>
    <span class="c1"># list of all strengths for color map</span>
    <span class="n">all_strengths</span> <span class="o">=</span> <span class="p">[]</span>
    <span class="c1"># Add attributes, contemporaneous and lagged links are handled separately</span>
    <span class="k">for</span> <span class="p">(</span><span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">,</span> <span class="n">dic</span><span class="p">)</span> <span class="ow">in</span> <span class="n">G</span><span class="o">.</span><span class="n">edges</span><span class="p">(</span><span class="n">data</span><span class="o">=</span><span class="kc">True</span><span class="p">):</span>
        <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;no_links&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">no_links</span>
        <span class="k">if</span> <span class="n">u</span> <span class="o">!=</span> <span class="n">v</span><span class="p">:</span>
            <span class="c1"># tau = np.abs((u - v) % max_lag)</span>
            <span class="c1"># Determine neighbors in TSG</span>
            <span class="n">i</span> <span class="o">=</span> <span class="n">u</span> <span class="o">//</span> <span class="n">max_lag</span>
            <span class="n">taui</span> <span class="o">=</span> <span class="o">-</span><span class="p">(</span><span class="n">max_lag</span> <span class="o">-</span><span class="mi">1</span> <span class="o">-</span> <span class="p">(</span><span class="n">u</span> <span class="o">%</span> <span class="n">max_lag</span><span class="p">))</span>
            <span class="n">j</span> <span class="o">=</span> <span class="n">v</span> <span class="o">//</span> <span class="n">max_lag</span>
            <span class="n">tauj</span> <span class="o">=</span> <span class="o">-</span><span class="p">(</span><span class="n">max_lag</span> <span class="o">-</span><span class="mi">1</span> <span class="o">-</span> <span class="p">(</span><span class="n">v</span> <span class="o">%</span> <span class="n">max_lag</span><span class="p">))</span>

            <span class="k">if</span> <span class="n">np</span><span class="o">.</span><span class="n">abs</span><span class="p">(</span><span class="n">i</span><span class="o">-</span><span class="n">j</span><span class="p">)</span> <span class="o">&lt;=</span> <span class="mi">1</span> <span class="ow">and</span> <span class="n">np</span><span class="o">.</span><span class="n">abs</span><span class="p">(</span><span class="n">tauj</span><span class="o">-</span><span class="n">taui</span><span class="p">)</span> <span class="o">&lt;=</span> <span class="mi">1</span><span class="p">:</span>
                <span class="n">inout</span> <span class="o">=</span> <span class="s1">&#39;inner&#39;</span>
                <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;inner_edge&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="kc">True</span>
                <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;outer_edge&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="kc">False</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">inout</span> <span class="o">=</span> <span class="s1">&#39;outer&#39;</span>
                <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;inner_edge&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="kc">False</span>
                <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;outer_edge&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="kc">True</span>

            <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;</span><span class="si">%s</span><span class="s2">_edge_type&quot;</span> <span class="o">%</span> <span class="n">inout</span><span class="p">]</span> <span class="o">=</span> <span class="n">tsg_style</span><span class="p">[</span><span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">]</span>

            <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;</span><span class="si">%s</span><span class="s2">_edge_alpha&quot;</span> <span class="o">%</span> <span class="n">inout</span><span class="p">]</span> <span class="o">=</span> <span class="n">alpha</span>

            <span class="k">if</span> <span class="n">link_width</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
                <span class="c1"># fraction of nonzero values</span>
                <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;</span><span class="si">%s</span><span class="s2">_edge_width&quot;</span> <span class="o">%</span> <span class="n">inout</span><span class="p">]</span> <span class="o">=</span> <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;</span><span class="si">%s</span><span class="s2">_edge_width&quot;</span> <span class="o">%</span> <span class="n">inout</span><span class="p">]</span> <span class="o">=</span> <span class="n">arrow_linewidth</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;</span><span class="si">%s</span><span class="s2">_edge_width&quot;</span> <span class="o">%</span> <span class="n">inout</span><span class="p">]</span> <span class="o">=</span> <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;</span><span class="si">%s</span><span class="s2">_edge_width&quot;</span> <span class="o">%</span> <span class="n">inout</span><span class="p">]</span> <span class="o">=</span> <span class="n">tsg_width</span><span class="p">[</span><span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">]</span>

            <span class="k">if</span> <span class="n">link_attribute</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
                <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;</span><span class="si">%s</span><span class="s2">_edge_attribute&quot;</span> <span class="o">%</span> <span class="n">inout</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;</span><span class="si">%s</span><span class="s2">_edge_attribute&quot;</span> <span class="o">%</span> <span class="n">inout</span><span class="p">]</span> <span class="o">=</span> <span class="n">tsg_attr</span><span class="p">[</span><span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">]</span>

            <span class="c1"># value at argmax of average</span>
            <span class="k">if</span> <span class="n">no_coloring</span><span class="p">:</span>
                <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;</span><span class="si">%s</span><span class="s2">_edge_color&quot;</span> <span class="o">%</span> <span class="n">inout</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;</span><span class="si">%s</span><span class="s2">_edge_color&quot;</span> <span class="o">%</span> <span class="n">inout</span><span class="p">]</span> <span class="o">=</span> <span class="n">tsg_val</span><span class="p">[</span><span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">]</span>

            <span class="n">all_strengths</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">dic</span><span class="p">[</span><span class="s2">&quot;</span><span class="si">%s</span><span class="s2">_edge_color&quot;</span> <span class="o">%</span> <span class="n">inout</span><span class="p">])</span>
            <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;label&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="c1"># print(u, v, dic)</span>

    <span class="c1"># If no links are present, set value to zero</span>
    <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">all_strengths</span><span class="p">)</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
        <span class="n">all_strengths</span> <span class="o">=</span> <span class="p">[</span><span class="mf">0.0</span><span class="p">]</span>

    <span class="n">posarray</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="n">N</span> <span class="o">*</span> <span class="n">max_lag</span><span class="p">,</span> <span class="mi">2</span><span class="p">))</span>
    <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">N</span> <span class="o">*</span> <span class="n">max_lag</span><span class="p">):</span>
        <span class="n">posarray</span><span class="p">[</span><span class="n">i</span><span class="p">]</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">([(</span><span class="n">i</span> <span class="o">%</span> <span class="n">max_lag</span><span class="p">),</span> <span class="p">(</span><span class="mf">1.0</span> <span class="o">-</span> <span class="n">i</span> <span class="o">//</span> <span class="n">max_lag</span><span class="p">)])</span>

    <span class="n">pos_tmp</span> <span class="o">=</span> <span class="p">{}</span>
    <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">N</span> <span class="o">*</span> <span class="n">max_lag</span><span class="p">):</span>
        <span class="c1"># for n in range(N):</span>
        <span class="c1">#     for tau in range(max_lag):</span>
        <span class="c1">#         i = n*N + tau</span>
        <span class="n">pos_tmp</span><span class="p">[</span><span class="n">i</span><span class="p">]</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">(</span>
            <span class="p">[</span>
                <span class="p">((</span><span class="n">i</span> <span class="o">%</span> <span class="n">max_lag</span><span class="p">)</span> <span class="o">-</span> <span class="n">posarray</span><span class="o">.</span><span class="n">min</span><span class="p">(</span><span class="n">axis</span><span class="o">=</span><span class="mi">0</span><span class="p">)[</span><span class="mi">0</span><span class="p">])</span>
                <span class="o">/</span> <span class="p">(</span><span class="n">posarray</span><span class="o">.</span><span class="n">max</span><span class="p">(</span><span class="n">axis</span><span class="o">=</span><span class="mi">0</span><span class="p">)[</span><span class="mi">0</span><span class="p">]</span> <span class="o">-</span> <span class="n">posarray</span><span class="o">.</span><span class="n">min</span><span class="p">(</span><span class="n">axis</span><span class="o">=</span><span class="mi">0</span><span class="p">)[</span><span class="mi">0</span><span class="p">]),</span>
                <span class="p">((</span><span class="mf">1.0</span> <span class="o">-</span> <span class="n">i</span> <span class="o">//</span> <span class="n">max_lag</span><span class="p">)</span> <span class="o">-</span> <span class="n">posarray</span><span class="o">.</span><span class="n">min</span><span class="p">(</span><span class="n">axis</span><span class="o">=</span><span class="mi">0</span><span class="p">)[</span><span class="mi">1</span><span class="p">])</span>
                <span class="o">/</span> <span class="p">(</span><span class="n">posarray</span><span class="o">.</span><span class="n">max</span><span class="p">(</span><span class="n">axis</span><span class="o">=</span><span class="mi">0</span><span class="p">)[</span><span class="mi">1</span><span class="p">]</span> <span class="o">-</span> <span class="n">posarray</span><span class="o">.</span><span class="n">min</span><span class="p">(</span><span class="n">axis</span><span class="o">=</span><span class="mi">0</span><span class="p">)[</span><span class="mi">1</span><span class="p">]),</span>
            <span class="p">]</span>
        <span class="p">)</span>
        <span class="n">pos_tmp</span><span class="p">[</span><span class="n">i</span><span class="p">][</span><span class="n">np</span><span class="o">.</span><span class="n">isnan</span><span class="p">(</span><span class="n">pos_tmp</span><span class="p">[</span><span class="n">i</span><span class="p">])]</span> <span class="o">=</span> <span class="mf">0.0</span>

    <span class="n">pos</span> <span class="o">=</span> <span class="p">{}</span>
    <span class="k">for</span> <span class="n">n</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">N</span><span class="p">):</span>
        <span class="k">for</span> <span class="n">tau</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">max_lag</span><span class="p">):</span>
            <span class="n">pos</span><span class="p">[</span><span class="n">n</span> <span class="o">*</span> <span class="n">max_lag</span> <span class="o">+</span> <span class="n">tau</span><span class="p">]</span> <span class="o">=</span> <span class="n">pos_tmp</span><span class="p">[</span><span class="n">order</span><span class="p">[</span><span class="n">n</span><span class="p">]</span> <span class="o">*</span> <span class="n">max_lag</span> <span class="o">+</span> <span class="n">tau</span><span class="p">]</span>

    <span class="n">node_rings</span> <span class="o">=</span> <span class="p">{</span>
        <span class="mi">0</span><span class="p">:</span> <span class="p">{</span><span class="s2">&quot;sizes&quot;</span><span class="p">:</span> <span class="kc">None</span><span class="p">,</span> <span class="s2">&quot;color_array&quot;</span><span class="p">:</span> <span class="kc">None</span><span class="p">,</span> <span class="s2">&quot;label&quot;</span><span class="p">:</span> <span class="s2">&quot;&quot;</span><span class="p">,</span> <span class="s2">&quot;colorbar&quot;</span><span class="p">:</span> <span class="kc">False</span><span class="p">,}</span>
    <span class="p">}</span>

    <span class="n">node_labels</span> <span class="o">=</span> <span class="p">[</span><span class="s2">&quot;&quot;</span> <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">N</span> <span class="o">*</span> <span class="n">max_lag</span><span class="p">)]</span>

    <span class="k">if</span> <span class="n">graph</span><span class="o">.</span><span class="n">ndim</span> <span class="o">==</span> <span class="mi">4</span> <span class="ow">and</span> <span class="n">val_matrix</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">show_colorbar</span> <span class="o">=</span> <span class="kc">False</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">show_colorbar</span> <span class="o">=</span> <span class="kc">True</span>

    <span class="n">_draw_network_with_curved_edges</span><span class="p">(</span>
        <span class="n">fig</span><span class="o">=</span><span class="n">fig</span><span class="p">,</span>
        <span class="n">ax</span><span class="o">=</span><span class="n">ax</span><span class="p">,</span>
        <span class="n">G</span><span class="o">=</span><span class="n">deepcopy</span><span class="p">(</span><span class="n">G</span><span class="p">),</span>
        <span class="n">pos</span><span class="o">=</span><span class="n">pos</span><span class="p">,</span>
        <span class="n">node_rings</span><span class="o">=</span><span class="n">node_rings</span><span class="p">,</span>
        <span class="n">node_labels</span><span class="o">=</span><span class="n">node_labels</span><span class="p">,</span>
        <span class="c1"># node_label_size=node_label_size,</span>
        <span class="n">node_alpha</span><span class="o">=</span><span class="n">alpha</span><span class="p">,</span>
        <span class="n">standard_size</span><span class="o">=</span><span class="n">node_size</span><span class="p">,</span>
        <span class="n">node_aspect</span><span class="o">=</span><span class="n">node_aspect</span><span class="p">,</span>
        <span class="n">standard_cmap</span><span class="o">=</span><span class="s2">&quot;OrRd&quot;</span><span class="p">,</span>
        <span class="n">standard_color_nodes</span><span class="o">=</span><span class="n">standard_color_nodes</span><span class="p">,</span>
        <span class="n">standard_color_links</span><span class="o">=</span><span class="n">standard_color_links</span><span class="p">,</span>
        <span class="n">log_sizes</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
        <span class="n">cmap_links</span><span class="o">=</span><span class="n">cmap_edges</span><span class="p">,</span>
        <span class="n">links_vmin</span><span class="o">=</span><span class="n">vmin_edges</span><span class="p">,</span>
        <span class="n">links_vmax</span><span class="o">=</span><span class="n">vmax_edges</span><span class="p">,</span>
        <span class="n">links_ticks</span><span class="o">=</span><span class="n">edge_ticks</span><span class="p">,</span>
        <span class="c1"># link_label_fontsize=link_label_fontsize,</span>
        <span class="n">arrowstyle</span><span class="o">=</span><span class="s2">&quot;simple&quot;</span><span class="p">,</span>
        <span class="n">arrowhead_size</span><span class="o">=</span><span class="n">arrowhead_size</span><span class="p">,</span>
        <span class="n">curved_radius</span><span class="o">=</span><span class="n">curved_radius</span><span class="p">,</span>
        <span class="n">label_fontsize</span><span class="o">=</span><span class="n">label_fontsize</span><span class="p">,</span>
        <span class="n">tick_label_size</span><span class="o">=</span><span class="n">tick_label_size</span><span class="p">,</span>
        <span class="n">label_fraction</span><span class="o">=</span><span class="mf">0.5</span><span class="p">,</span>
        <span class="n">link_colorbar_label</span><span class="o">=</span><span class="n">link_colorbar_label</span><span class="p">,</span>
        <span class="n">inner_edge_curved</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
        <span class="c1"># network_lower_bound=network_lower_bound,</span>
        <span class="c1"># network_left_bound=label_space_left,</span>
        <span class="n">inner_edge_style</span><span class="o">=</span><span class="n">inner_edge_style</span><span class="p">,</span>
        <span class="n">special_nodes</span><span class="o">=</span><span class="n">special_nodes</span><span class="p">,</span>
        <span class="n">show_colorbar</span><span class="o">=</span><span class="n">show_colorbar</span><span class="p">,</span>
        <span class="n">node_classification</span><span class="o">=</span><span class="n">node_classification_tsg</span><span class="p">,</span>
        <span class="n">max_lag</span><span class="o">=</span><span class="n">max_lag</span><span class="p">,</span>
    <span class="p">)</span>

    <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">N</span><span class="p">):</span>
        <span class="n">trans</span> <span class="o">=</span> <span class="n">transforms</span><span class="o">.</span><span class="n">blended_transform_factory</span><span class="p">(</span><span class="n">ax</span><span class="o">.</span><span class="n">transAxes</span><span class="p">,</span> <span class="n">ax</span><span class="o">.</span><span class="n">transData</span><span class="p">)</span>
        <span class="c1"># trans = transforms.blended_transform_factory(fig.transFigure, ax.transData)</span>
        <span class="n">ax</span><span class="o">.</span><span class="n">text</span><span class="p">(</span>
            <span class="mf">0.</span><span class="p">,</span>
            <span class="n">pos</span><span class="p">[</span><span class="n">order</span><span class="p">[</span><span class="n">i</span><span class="p">]</span> <span class="o">*</span> <span class="n">max_lag</span><span class="p">][</span><span class="mi">1</span><span class="p">],</span>
            <span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">var_names</span><span class="p">[</span><span class="n">order</span><span class="p">[</span><span class="n">i</span><span class="p">]]</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">,</span>
            <span class="n">fontsize</span><span class="o">=</span><span class="n">label_fontsize</span><span class="p">,</span>
            <span class="n">horizontalalignment</span><span class="o">=</span><span class="s2">&quot;right&quot;</span><span class="p">,</span>
            <span class="n">verticalalignment</span><span class="o">=</span><span class="s2">&quot;center&quot;</span><span class="p">,</span>
            <span class="n">transform</span><span class="o">=</span><span class="n">trans</span><span class="p">,</span>
        <span class="p">)</span>

    <span class="k">for</span> <span class="n">tau</span> <span class="ow">in</span> <span class="n">np</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span><span class="n">max_lag</span> <span class="o">-</span> <span class="mi">1</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">):</span>
        <span class="n">trans</span> <span class="o">=</span> <span class="n">transforms</span><span class="o">.</span><span class="n">blended_transform_factory</span><span class="p">(</span><span class="n">ax</span><span class="o">.</span><span class="n">transData</span><span class="p">,</span> <span class="n">ax</span><span class="o">.</span><span class="n">transAxes</span><span class="p">)</span>
        <span class="c1"># trans = transforms.blended_transform_factory(ax.transData, fig.transFigure)</span>
        <span class="k">if</span> <span class="n">tau</span> <span class="o">==</span> <span class="n">max_lag</span> <span class="o">-</span> <span class="mi">1</span><span class="p">:</span>
            <span class="n">ax</span><span class="o">.</span><span class="n">text</span><span class="p">(</span>
                <span class="n">pos</span><span class="p">[</span><span class="n">tau</span><span class="p">][</span><span class="mi">0</span><span class="p">],</span>
                <span class="mf">1.0</span><span class="p">,</span> <span class="c1"># - label_space_top,</span>
                <span class="sa">r</span><span class="s2">&quot;$t$&quot;</span><span class="p">,</span>
                <span class="n">fontsize</span><span class="o">=</span><span class="nb">int</span><span class="p">(</span><span class="n">label_fontsize</span> <span class="o">*</span> <span class="mf">0.8</span><span class="p">),</span>
                <span class="n">horizontalalignment</span><span class="o">=</span><span class="s2">&quot;center&quot;</span><span class="p">,</span>
                <span class="n">verticalalignment</span><span class="o">=</span><span class="s2">&quot;bottom&quot;</span><span class="p">,</span>
                <span class="n">transform</span><span class="o">=</span><span class="n">trans</span><span class="p">,</span>
            <span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">ax</span><span class="o">.</span><span class="n">text</span><span class="p">(</span>
                <span class="n">pos</span><span class="p">[</span><span class="n">tau</span><span class="p">][</span><span class="mi">0</span><span class="p">],</span>
                <span class="mf">1.0</span><span class="p">,</span> <span class="c1"># - label_space_top,</span>
                <span class="sa">r</span><span class="s2">&quot;$t-</span><span class="si">%s</span><span class="s2">$&quot;</span> <span class="o">%</span> <span class="nb">str</span><span class="p">(</span><span class="n">max_lag</span> <span class="o">-</span> <span class="n">tau</span> <span class="o">-</span> <span class="mi">1</span><span class="p">),</span>
                <span class="n">fontsize</span><span class="o">=</span><span class="nb">int</span><span class="p">(</span><span class="n">label_fontsize</span> <span class="o">*</span> <span class="mf">0.8</span><span class="p">),</span>
                <span class="n">horizontalalignment</span><span class="o">=</span><span class="s2">&quot;center&quot;</span><span class="p">,</span>
                <span class="n">verticalalignment</span><span class="o">=</span><span class="s2">&quot;bottom&quot;</span><span class="p">,</span>
                <span class="n">transform</span><span class="o">=</span><span class="n">trans</span><span class="p">,</span>
            <span class="p">)</span>

    <span class="c1"># pyplot.tight_layout()</span>
    <span class="k">if</span> <span class="n">save_name</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">pyplot</span><span class="o">.</span><span class="n">savefig</span><span class="p">(</span><span class="n">save_name</span><span class="p">,</span> <span class="n">dpi</span><span class="o">=</span><span class="mi">300</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">fig</span><span class="p">,</span> <span class="n">ax</span></div>


<div class="viewcode-block" id="plot_mediation_time_series_graph"><a class="viewcode-back" href="../../index.html#tigramite.plotting.plot_mediation_time_series_graph">[docs]</a><span class="k">def</span> <span class="nf">plot_mediation_time_series_graph</span><span class="p">(</span>
    <span class="n">path_node_array</span><span class="p">,</span>
    <span class="n">tsg_path_val_matrix</span><span class="p">,</span>
    <span class="n">var_names</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
    <span class="n">fig_ax</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
    <span class="n">figsize</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
    <span class="n">link_colorbar_label</span><span class="o">=</span><span class="s2">&quot;link coeff. (edge color)&quot;</span><span class="p">,</span>
    <span class="n">node_colorbar_label</span><span class="o">=</span><span class="s2">&quot;MCE (node color)&quot;</span><span class="p">,</span>
    <span class="n">save_name</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
    <span class="n">link_width</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
    <span class="n">arrow_linewidth</span><span class="o">=</span><span class="mi">8</span><span class="p">,</span>
    <span class="n">vmin_edges</span><span class="o">=-</span><span class="mi">1</span><span class="p">,</span>
    <span class="n">vmax_edges</span><span class="o">=</span><span class="mf">1.0</span><span class="p">,</span>
    <span class="n">edge_ticks</span><span class="o">=</span><span class="mf">0.4</span><span class="p">,</span>
    <span class="n">cmap_edges</span><span class="o">=</span><span class="s2">&quot;RdBu_r&quot;</span><span class="p">,</span>
    <span class="n">order</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
    <span class="n">vmin_nodes</span><span class="o">=-</span><span class="mf">1.0</span><span class="p">,</span>
    <span class="n">vmax_nodes</span><span class="o">=</span><span class="mf">1.0</span><span class="p">,</span>
    <span class="n">node_ticks</span><span class="o">=</span><span class="mf">0.4</span><span class="p">,</span>
    <span class="n">cmap_nodes</span><span class="o">=</span><span class="s2">&quot;RdBu_r&quot;</span><span class="p">,</span>
    <span class="n">node_size</span><span class="o">=</span><span class="mf">0.1</span><span class="p">,</span>
    <span class="n">node_aspect</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
    <span class="n">arrowhead_size</span><span class="o">=</span><span class="mi">20</span><span class="p">,</span>
    <span class="n">curved_radius</span><span class="o">=</span><span class="mf">0.2</span><span class="p">,</span>
    <span class="n">label_fontsize</span><span class="o">=</span><span class="mi">12</span><span class="p">,</span>
    <span class="n">alpha</span><span class="o">=</span><span class="mf">1.0</span><span class="p">,</span>
    <span class="n">node_label_size</span><span class="o">=</span><span class="mi">12</span><span class="p">,</span>
    <span class="n">tick_label_size</span><span class="o">=</span><span class="mi">6</span><span class="p">,</span>
    <span class="n">standard_color_links</span><span class="o">=</span><span class="s1">&#39;black&#39;</span><span class="p">,</span>
    <span class="n">standard_color_nodes</span><span class="o">=</span><span class="s1">&#39;lightgrey&#39;</span><span class="p">,</span>
<span class="p">):</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;Creates a mediation time series graph plot.</span>
<span class="sd">    This is still in beta. The time series graph&#39;s links are colored by</span>
<span class="sd">    val_matrix.</span>

<span class="sd">    Parameters</span>
<span class="sd">    ----------</span>
<span class="sd">    tsg_path_val_matrix : array_like</span>
<span class="sd">        Matrix of shape (N*tau_max, N*tau_max) containing link weight values.</span>
<span class="sd">    path_node_array: array_like</span>
<span class="sd">        Array of shape (N,) containing node values.</span>
<span class="sd">    var_names : list, optional (default: None)</span>
<span class="sd">        List of variable names. If None, range(N) is used.</span>
<span class="sd">    fig_ax : tuple of figure and axis object, optional (default: None)</span>
<span class="sd">        Figure and axes instance. If None they are created.</span>
<span class="sd">    figsize : tuple</span>
<span class="sd">        Size of figure.</span>
<span class="sd">    save_name : str, optional (default: None)</span>
<span class="sd">        Name of figure file to save figure. If None, figure is shown in window.</span>
<span class="sd">    link_colorbar_label : str, optional (default: &#39;link coeff. (edge color)&#39;)</span>
<span class="sd">        Link colorbar label.</span>
<span class="sd">    node_colorbar_label : str, optional (default: &#39;MCE (node color)&#39;)</span>
<span class="sd">        Node colorbar label.</span>
<span class="sd">    link_width : array-like, optional (default: None)</span>
<span class="sd">        Array of val_matrix.shape specifying relative link width with maximum</span>
<span class="sd">        given by arrow_linewidth. If None, all links have same width.</span>
<span class="sd">    order : list, optional (default: None)</span>
<span class="sd">        order of variables from top to bottom.</span>
<span class="sd">    arrow_linewidth : float, optional (default: 30)</span>
<span class="sd">        Linewidth.</span>
<span class="sd">    vmin_edges : float, optional (default: -1)</span>
<span class="sd">        Link colorbar scale lower bound.</span>
<span class="sd">    vmax_edges : float, optional (default: 1)</span>
<span class="sd">        Link colorbar scale upper bound.</span>
<span class="sd">    edge_ticks : float, optional (default: 0.4)</span>
<span class="sd">        Link tick mark interval.</span>
<span class="sd">    cmap_edges : str, optional (default: &#39;RdBu_r&#39;)</span>
<span class="sd">        Colormap for links.</span>
<span class="sd">    vmin_nodes : float, optional (default: 0)</span>
<span class="sd">        Node colorbar scale lower bound.</span>
<span class="sd">    vmax_nodes : float, optional (default: 1)</span>
<span class="sd">        Node colorbar scale upper bound.</span>
<span class="sd">    node_ticks : float, optional (default: 0.4)</span>
<span class="sd">        Node tick mark interval.</span>
<span class="sd">    cmap_nodes : str, optional (default: &#39;OrRd&#39;)</span>
<span class="sd">        Colormap for links.</span>
<span class="sd">    node_size : int, optional (default: 0.1)</span>
<span class="sd">        Node size.</span>
<span class="sd">    node_aspect : float, optional (default: None)</span>
<span class="sd">        Ratio between the heigth and width of the varible nodes.</span>
<span class="sd">    arrowhead_size : int, optional (default: 20)</span>
<span class="sd">        Size of link arrow head. Passed on to FancyArrowPatch object.</span>
<span class="sd">    curved_radius, float, optional (default: 0.2)</span>
<span class="sd">        Curvature of links. Passed on to FancyArrowPatch object.</span>
<span class="sd">    label_fontsize : int, optional (default: 10)</span>
<span class="sd">        Fontsize of colorbar labels.</span>
<span class="sd">    alpha : float, optional (default: 1.)</span>
<span class="sd">        Opacity.</span>
<span class="sd">    node_label_size : int, optional (default: 10)</span>
<span class="sd">        Fontsize of node labels.</span>
<span class="sd">    link_label_fontsize : int, optional (default: 6)</span>
<span class="sd">        Fontsize of link labels.</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">N</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="n">path_node_array</span><span class="p">)</span>
    <span class="n">Nmaxlag</span> <span class="o">=</span> <span class="n">tsg_path_val_matrix</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
    <span class="n">max_lag</span> <span class="o">=</span> <span class="n">Nmaxlag</span> <span class="o">//</span> <span class="n">N</span>

    <span class="k">if</span> <span class="n">var_names</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">var_names</span> <span class="o">=</span> <span class="nb">range</span><span class="p">(</span><span class="n">N</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">fig_ax</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">fig</span> <span class="o">=</span> <span class="n">pyplot</span><span class="o">.</span><span class="n">figure</span><span class="p">(</span><span class="n">figsize</span><span class="o">=</span><span class="n">figsize</span><span class="p">)</span>
        <span class="n">ax</span> <span class="o">=</span> <span class="n">fig</span><span class="o">.</span><span class="n">add_subplot</span><span class="p">(</span><span class="mi">111</span><span class="p">,</span> <span class="n">frame_on</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">fig</span><span class="p">,</span> <span class="n">ax</span> <span class="o">=</span> <span class="n">fig_ax</span>

    <span class="k">if</span> <span class="n">link_width</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> <span class="ow">not</span> <span class="n">np</span><span class="o">.</span><span class="n">all</span><span class="p">(</span><span class="n">link_width</span> <span class="o">&gt;=</span> <span class="mf">0.0</span><span class="p">):</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;link_width must be non-negative&quot;</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">order</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">order</span> <span class="o">=</span> <span class="nb">range</span><span class="p">(</span><span class="n">N</span><span class="p">)</span>

    <span class="k">if</span> <span class="nb">set</span><span class="p">(</span><span class="n">order</span><span class="p">)</span> <span class="o">!=</span> <span class="nb">set</span><span class="p">(</span><span class="nb">range</span><span class="p">(</span><span class="n">N</span><span class="p">)):</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;order must be a permutation of range(N)&quot;</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">translate</span><span class="p">(</span><span class="n">row</span><span class="p">,</span> <span class="n">lag</span><span class="p">):</span>
        <span class="k">return</span> <span class="n">row</span> <span class="o">*</span> <span class="n">max_lag</span> <span class="o">+</span> <span class="n">lag</span>

    <span class="k">if</span> <span class="n">np</span><span class="o">.</span><span class="n">count_nonzero</span><span class="p">(</span><span class="n">tsg_path_val_matrix</span><span class="p">)</span> <span class="o">==</span> <span class="n">np</span><span class="o">.</span><span class="n">count_nonzero</span><span class="p">(</span>
        <span class="n">np</span><span class="o">.</span><span class="n">diagonal</span><span class="p">(</span><span class="n">tsg_path_val_matrix</span><span class="p">)</span>
    <span class="p">):</span>
        <span class="n">diagonal</span> <span class="o">=</span> <span class="kc">True</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">diagonal</span> <span class="o">=</span> <span class="kc">False</span>

    <span class="k">if</span> <span class="n">np</span><span class="o">.</span><span class="n">count_nonzero</span><span class="p">(</span><span class="n">tsg_path_val_matrix</span><span class="p">)</span> <span class="o">==</span> <span class="n">tsg_path_val_matrix</span><span class="o">.</span><span class="n">size</span> <span class="ow">or</span> <span class="n">diagonal</span><span class="p">:</span>
        <span class="n">tsg_path_val_matrix</span><span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">]</span> <span class="o">=</span> <span class="mi">1</span>
        <span class="n">no_links</span> <span class="o">=</span> <span class="kc">True</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">no_links</span> <span class="o">=</span> <span class="kc">False</span>

    <span class="c1"># Define graph links by absolute maximum (positive or negative like for</span>
    <span class="c1"># partial correlation)</span>
    <span class="n">tsg</span> <span class="o">=</span> <span class="n">tsg_path_val_matrix</span>
    <span class="n">tsg_attr</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="n">N</span> <span class="o">*</span> <span class="n">max_lag</span><span class="p">,</span> <span class="n">N</span> <span class="o">*</span> <span class="n">max_lag</span><span class="p">))</span>

    <span class="n">G</span> <span class="o">=</span> <span class="n">nx</span><span class="o">.</span><span class="n">DiGraph</span><span class="p">(</span><span class="n">tsg</span><span class="p">)</span>

    <span class="c1"># node_color = np.zeros(N)</span>
    <span class="c1"># list of all strengths for color map</span>
    <span class="n">all_strengths</span> <span class="o">=</span> <span class="p">[]</span>
    <span class="c1"># Add attributes, contemporaneous and lagged links are handled separately</span>
    <span class="k">for</span> <span class="p">(</span><span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">,</span> <span class="n">dic</span><span class="p">)</span> <span class="ow">in</span> <span class="n">G</span><span class="o">.</span><span class="n">edges</span><span class="p">(</span><span class="n">data</span><span class="o">=</span><span class="kc">True</span><span class="p">):</span>
        <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;no_links&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">no_links</span>
        <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;outer_edge_attribute&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span>

        <span class="k">if</span> <span class="n">u</span> <span class="o">!=</span> <span class="n">v</span><span class="p">:</span>

            <span class="k">if</span> <span class="n">u</span> <span class="o">%</span> <span class="n">max_lag</span> <span class="o">==</span> <span class="n">v</span> <span class="o">%</span> <span class="n">max_lag</span><span class="p">:</span>
                <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;inner_edge&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="kc">True</span>
                <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;outer_edge&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="kc">False</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;inner_edge&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="kc">False</span>
                <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;outer_edge&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="kc">True</span>

            <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;inner_edge_alpha&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">alpha</span>
            <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;inner_edge_color&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">_get_absmax</span><span class="p">(</span>
                <span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">([[[</span><span class="n">tsg</span><span class="p">[</span><span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">],</span> <span class="n">tsg</span><span class="p">[</span><span class="n">v</span><span class="p">,</span> <span class="n">u</span><span class="p">]]]])</span>
            <span class="p">)</span><span class="o">.</span><span class="n">squeeze</span><span class="p">()</span>
            <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;inner_edge_width&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">arrow_linewidth</span>
            <span class="n">all_strengths</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">dic</span><span class="p">[</span><span class="s2">&quot;inner_edge_color&quot;</span><span class="p">])</span>

            <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;outer_edge_alpha&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">alpha</span>

            <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;outer_edge_width&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">arrow_linewidth</span>

            <span class="c1"># value at argmax of average</span>
            <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;outer_edge_color&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">tsg</span><span class="p">[</span><span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">]</span>
            <span class="n">all_strengths</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">dic</span><span class="p">[</span><span class="s2">&quot;outer_edge_color&quot;</span><span class="p">])</span>
            <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;label&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span>

        <span class="c1"># dic[&#39;outer_edge_edge&#39;] = False</span>
        <span class="c1"># dic[&#39;outer_edge_edgecolor&#39;] = None</span>
        <span class="c1"># dic[&#39;inner_edge_edge&#39;] = False</span>
        <span class="c1"># dic[&#39;inner_edge_edgecolor&#39;] = None</span>

    <span class="c1"># If no links are present, set value to zero</span>
    <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">all_strengths</span><span class="p">)</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
        <span class="n">all_strengths</span> <span class="o">=</span> <span class="p">[</span><span class="mf">0.0</span><span class="p">]</span>

    <span class="n">posarray</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="n">N</span> <span class="o">*</span> <span class="n">max_lag</span><span class="p">,</span> <span class="mi">2</span><span class="p">))</span>
    <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">N</span> <span class="o">*</span> <span class="n">max_lag</span><span class="p">):</span>
        <span class="n">posarray</span><span class="p">[</span><span class="n">i</span><span class="p">]</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">([(</span><span class="n">i</span> <span class="o">%</span> <span class="n">max_lag</span><span class="p">),</span> <span class="p">(</span><span class="mf">1.0</span> <span class="o">-</span> <span class="n">i</span> <span class="o">//</span> <span class="n">max_lag</span><span class="p">)])</span>

    <span class="n">pos_tmp</span> <span class="o">=</span> <span class="p">{}</span>
    <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">N</span> <span class="o">*</span> <span class="n">max_lag</span><span class="p">):</span>
        <span class="c1"># for n in range(N):</span>
        <span class="c1">#     for tau in range(max_lag):</span>
        <span class="c1">#         i = n*N + tau</span>
        <span class="n">pos_tmp</span><span class="p">[</span><span class="n">i</span><span class="p">]</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">(</span>
            <span class="p">[</span>
                <span class="p">((</span><span class="n">i</span> <span class="o">%</span> <span class="n">max_lag</span><span class="p">)</span> <span class="o">-</span> <span class="n">posarray</span><span class="o">.</span><span class="n">min</span><span class="p">(</span><span class="n">axis</span><span class="o">=</span><span class="mi">0</span><span class="p">)[</span><span class="mi">0</span><span class="p">])</span>
                <span class="o">/</span> <span class="p">(</span><span class="n">posarray</span><span class="o">.</span><span class="n">max</span><span class="p">(</span><span class="n">axis</span><span class="o">=</span><span class="mi">0</span><span class="p">)[</span><span class="mi">0</span><span class="p">]</span> <span class="o">-</span> <span class="n">posarray</span><span class="o">.</span><span class="n">min</span><span class="p">(</span><span class="n">axis</span><span class="o">=</span><span class="mi">0</span><span class="p">)[</span><span class="mi">0</span><span class="p">]),</span>
                <span class="p">((</span><span class="mf">1.0</span> <span class="o">-</span> <span class="n">i</span> <span class="o">//</span> <span class="n">max_lag</span><span class="p">)</span> <span class="o">-</span> <span class="n">posarray</span><span class="o">.</span><span class="n">min</span><span class="p">(</span><span class="n">axis</span><span class="o">=</span><span class="mi">0</span><span class="p">)[</span><span class="mi">1</span><span class="p">])</span>
                <span class="o">/</span> <span class="p">(</span><span class="n">posarray</span><span class="o">.</span><span class="n">max</span><span class="p">(</span><span class="n">axis</span><span class="o">=</span><span class="mi">0</span><span class="p">)[</span><span class="mi">1</span><span class="p">]</span> <span class="o">-</span> <span class="n">posarray</span><span class="o">.</span><span class="n">min</span><span class="p">(</span><span class="n">axis</span><span class="o">=</span><span class="mi">0</span><span class="p">)[</span><span class="mi">1</span><span class="p">]),</span>
            <span class="p">]</span>
        <span class="p">)</span>
        <span class="n">pos_tmp</span><span class="p">[</span><span class="n">i</span><span class="p">][</span><span class="n">np</span><span class="o">.</span><span class="n">isnan</span><span class="p">(</span><span class="n">pos_tmp</span><span class="p">[</span><span class="n">i</span><span class="p">])]</span> <span class="o">=</span> <span class="mf">0.0</span>

    <span class="n">pos</span> <span class="o">=</span> <span class="p">{}</span>
    <span class="k">for</span> <span class="n">n</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">N</span><span class="p">):</span>
        <span class="k">for</span> <span class="n">tau</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">max_lag</span><span class="p">):</span>
            <span class="n">pos</span><span class="p">[</span><span class="n">n</span> <span class="o">*</span> <span class="n">max_lag</span> <span class="o">+</span> <span class="n">tau</span><span class="p">]</span> <span class="o">=</span> <span class="n">pos_tmp</span><span class="p">[</span><span class="n">order</span><span class="p">[</span><span class="n">n</span><span class="p">]</span> <span class="o">*</span> <span class="n">max_lag</span> <span class="o">+</span> <span class="n">tau</span><span class="p">]</span>

    <span class="n">node_color</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="n">N</span> <span class="o">*</span> <span class="n">max_lag</span><span class="p">)</span>
    <span class="k">for</span> <span class="n">inet</span><span class="p">,</span> <span class="n">n</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="nb">range</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="n">N</span> <span class="o">*</span> <span class="n">max_lag</span><span class="p">,</span> <span class="n">max_lag</span><span class="p">)):</span>
        <span class="n">node_color</span><span class="p">[</span><span class="n">n</span> <span class="p">:</span> <span class="n">n</span> <span class="o">+</span> <span class="n">max_lag</span><span class="p">]</span> <span class="o">=</span> <span class="n">path_node_array</span><span class="p">[</span><span class="n">inet</span><span class="p">]</span>

    <span class="c1"># node_rings = {0: {&#39;sizes&#39;: None, &#39;color_array&#39;: color_array,</span>
    <span class="c1">#                   &#39;label&#39;: &#39;&#39;, &#39;colorbar&#39;: False,</span>
    <span class="c1">#                   }</span>
    <span class="c1">#               }</span>

    <span class="n">node_rings</span> <span class="o">=</span> <span class="p">{</span>
        <span class="mi">0</span><span class="p">:</span> <span class="p">{</span>
            <span class="s2">&quot;sizes&quot;</span><span class="p">:</span> <span class="kc">None</span><span class="p">,</span>
            <span class="s2">&quot;color_array&quot;</span><span class="p">:</span> <span class="n">node_color</span><span class="p">,</span>
            <span class="s2">&quot;cmap&quot;</span><span class="p">:</span> <span class="n">cmap_nodes</span><span class="p">,</span>
            <span class="s2">&quot;vmin&quot;</span><span class="p">:</span> <span class="n">vmin_nodes</span><span class="p">,</span>
            <span class="s2">&quot;vmax&quot;</span><span class="p">:</span> <span class="n">vmax_nodes</span><span class="p">,</span>
            <span class="s2">&quot;ticks&quot;</span><span class="p">:</span> <span class="n">node_ticks</span><span class="p">,</span>
            <span class="s2">&quot;label&quot;</span><span class="p">:</span> <span class="n">node_colorbar_label</span><span class="p">,</span>
            <span class="s2">&quot;colorbar&quot;</span><span class="p">:</span> <span class="kc">True</span><span class="p">,</span>
        <span class="p">}</span>
    <span class="p">}</span>

    <span class="c1"># ] for v in range(max_lag)]</span>
    <span class="n">node_labels</span> <span class="o">=</span> <span class="p">[</span><span class="s2">&quot;&quot;</span> <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">N</span> <span class="o">*</span> <span class="n">max_lag</span><span class="p">)]</span>

    <span class="n">_draw_network_with_curved_edges</span><span class="p">(</span>
        <span class="n">fig</span><span class="o">=</span><span class="n">fig</span><span class="p">,</span>
        <span class="n">ax</span><span class="o">=</span><span class="n">ax</span><span class="p">,</span>
        <span class="n">G</span><span class="o">=</span><span class="n">deepcopy</span><span class="p">(</span><span class="n">G</span><span class="p">),</span>
        <span class="n">pos</span><span class="o">=</span><span class="n">pos</span><span class="p">,</span>
        <span class="c1"># dictionary of rings: {0:{&#39;sizes&#39;:(N,)-array, &#39;color_array&#39;:(N,)-array</span>
        <span class="c1"># or None, &#39;cmap&#39;:string,</span>
        <span class="n">node_rings</span><span class="o">=</span><span class="n">node_rings</span><span class="p">,</span>
        <span class="c1"># &#39;vmin&#39;:float or None, &#39;vmax&#39;:float or None, &#39;label&#39;:string or None}}</span>
        <span class="n">node_labels</span><span class="o">=</span><span class="n">node_labels</span><span class="p">,</span>
        <span class="n">node_label_size</span><span class="o">=</span><span class="n">node_label_size</span><span class="p">,</span>
        <span class="n">node_alpha</span><span class="o">=</span><span class="n">alpha</span><span class="p">,</span>
        <span class="n">standard_size</span><span class="o">=</span><span class="n">node_size</span><span class="p">,</span>
        <span class="n">node_aspect</span><span class="o">=</span><span class="n">node_aspect</span><span class="p">,</span>
        <span class="n">standard_cmap</span><span class="o">=</span><span class="s2">&quot;OrRd&quot;</span><span class="p">,</span>
        <span class="n">standard_color_nodes</span><span class="o">=</span><span class="n">standard_color_nodes</span><span class="p">,</span>
        <span class="n">standard_color_links</span><span class="o">=</span><span class="n">standard_color_links</span><span class="p">,</span>
        <span class="n">log_sizes</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
        <span class="n">cmap_links</span><span class="o">=</span><span class="n">cmap_edges</span><span class="p">,</span>
        <span class="n">links_vmin</span><span class="o">=</span><span class="n">vmin_edges</span><span class="p">,</span>
        <span class="n">links_vmax</span><span class="o">=</span><span class="n">vmax_edges</span><span class="p">,</span>
        <span class="n">links_ticks</span><span class="o">=</span><span class="n">edge_ticks</span><span class="p">,</span>
        <span class="n">tick_label_size</span><span class="o">=</span><span class="n">tick_label_size</span><span class="p">,</span>
        <span class="c1"># cmap_links_edges=&#39;YlOrRd&#39;, links_edges_vmin=-1., links_edges_vmax=1.,</span>
        <span class="c1"># links_edges_ticks=.2, link_edge_colorbar_label=&#39;link_edge&#39;,</span>
        <span class="n">arrowhead_size</span><span class="o">=</span><span class="n">arrowhead_size</span><span class="p">,</span>
        <span class="n">curved_radius</span><span class="o">=</span><span class="n">curved_radius</span><span class="p">,</span>
        <span class="n">label_fontsize</span><span class="o">=</span><span class="n">label_fontsize</span><span class="p">,</span>
        <span class="n">label_fraction</span><span class="o">=</span><span class="mf">0.5</span><span class="p">,</span>
        <span class="n">link_colorbar_label</span><span class="o">=</span><span class="n">link_colorbar_label</span><span class="p">,</span>
        <span class="n">inner_edge_curved</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
        <span class="c1"># network_lower_bound=network_lower_bound</span>
        <span class="c1"># inner_edge_style=inner_edge_style</span>
    <span class="p">)</span>

    <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">N</span><span class="p">):</span>
        <span class="n">trans</span> <span class="o">=</span> <span class="n">transforms</span><span class="o">.</span><span class="n">blended_transform_factory</span><span class="p">(</span><span class="n">ax</span><span class="o">.</span><span class="n">transAxes</span><span class="p">,</span> <span class="n">ax</span><span class="o">.</span><span class="n">transData</span><span class="p">)</span>
        <span class="c1"># trans = transforms.blended_transform_factory(fig.transFigure, ax.transData)</span>
        <span class="n">ax</span><span class="o">.</span><span class="n">text</span><span class="p">(</span>
            <span class="mf">0.</span><span class="p">,</span>
            <span class="n">pos</span><span class="p">[</span><span class="n">order</span><span class="p">[</span><span class="n">i</span><span class="p">]</span> <span class="o">*</span> <span class="n">max_lag</span><span class="p">][</span><span class="mi">1</span><span class="p">],</span>
            <span class="s2">&quot;</span><span class="si">%s</span><span class="s2">&quot;</span> <span class="o">%</span> <span class="nb">str</span><span class="p">(</span><span class="n">var_names</span><span class="p">[</span><span class="n">order</span><span class="p">[</span><span class="n">i</span><span class="p">]]),</span>
            <span class="n">fontsize</span><span class="o">=</span><span class="n">label_fontsize</span><span class="p">,</span>
            <span class="n">horizontalalignment</span><span class="o">=</span><span class="s2">&quot;right&quot;</span><span class="p">,</span>
            <span class="n">verticalalignment</span><span class="o">=</span><span class="s2">&quot;center&quot;</span><span class="p">,</span>
            <span class="n">transform</span><span class="o">=</span><span class="n">trans</span><span class="p">,</span>
        <span class="p">)</span>

    <span class="k">for</span> <span class="n">tau</span> <span class="ow">in</span> <span class="n">np</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span><span class="n">max_lag</span> <span class="o">-</span> <span class="mi">1</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">):</span>
        <span class="n">trans</span> <span class="o">=</span> <span class="n">transforms</span><span class="o">.</span><span class="n">blended_transform_factory</span><span class="p">(</span><span class="n">ax</span><span class="o">.</span><span class="n">transData</span><span class="p">,</span> <span class="n">ax</span><span class="o">.</span><span class="n">transAxes</span><span class="p">)</span>
        <span class="c1"># trans = transforms.blended_transform_factory(ax.transData, fig.transFigure)</span>
        <span class="k">if</span> <span class="n">tau</span> <span class="o">==</span> <span class="n">max_lag</span> <span class="o">-</span> <span class="mi">1</span><span class="p">:</span>
            <span class="n">ax</span><span class="o">.</span><span class="n">text</span><span class="p">(</span>
                <span class="n">pos</span><span class="p">[</span><span class="n">tau</span><span class="p">][</span><span class="mi">0</span><span class="p">],</span>
                <span class="mf">1.0</span><span class="p">,</span> <span class="c1"># - label_space_top,</span>
                <span class="sa">r</span><span class="s2">&quot;$t$&quot;</span><span class="p">,</span>
                <span class="n">fontsize</span><span class="o">=</span><span class="n">label_fontsize</span><span class="p">,</span>
                <span class="n">horizontalalignment</span><span class="o">=</span><span class="s2">&quot;center&quot;</span><span class="p">,</span>
                <span class="n">verticalalignment</span><span class="o">=</span><span class="s2">&quot;bottom&quot;</span><span class="p">,</span>
                <span class="n">transform</span><span class="o">=</span><span class="n">trans</span><span class="p">,</span>
            <span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">ax</span><span class="o">.</span><span class="n">text</span><span class="p">(</span>
                <span class="n">pos</span><span class="p">[</span><span class="n">tau</span><span class="p">][</span><span class="mi">0</span><span class="p">],</span>
                <span class="mf">1.0</span><span class="p">,</span> <span class="c1"># - label_space_top,</span>
                <span class="sa">r</span><span class="s2">&quot;$t-</span><span class="si">%s</span><span class="s2">$&quot;</span> <span class="o">%</span> <span class="nb">str</span><span class="p">(</span><span class="n">max_lag</span> <span class="o">-</span> <span class="n">tau</span> <span class="o">-</span> <span class="mi">1</span><span class="p">),</span>
                <span class="n">fontsize</span><span class="o">=</span><span class="n">label_fontsize</span><span class="p">,</span>
                <span class="n">horizontalalignment</span><span class="o">=</span><span class="s2">&quot;center&quot;</span><span class="p">,</span>
                <span class="n">verticalalignment</span><span class="o">=</span><span class="s2">&quot;bottom&quot;</span><span class="p">,</span>
                <span class="n">transform</span><span class="o">=</span><span class="n">trans</span><span class="p">,</span>
            <span class="p">)</span>

    <span class="c1"># fig.subplots_adjust(left=0.1, right=.98, bottom=.25, top=.9)</span>
    <span class="c1"># savestring = os.path.expanduser(save_name)</span>
    <span class="k">if</span> <span class="n">save_name</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">pyplot</span><span class="o">.</span><span class="n">savefig</span><span class="p">(</span><span class="n">save_name</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">pyplot</span><span class="o">.</span><span class="n">show</span><span class="p">()</span></div>


<div class="viewcode-block" id="plot_mediation_graph"><a class="viewcode-back" href="../../index.html#tigramite.plotting.plot_mediation_graph">[docs]</a><span class="k">def</span> <span class="nf">plot_mediation_graph</span><span class="p">(</span>
    <span class="n">path_val_matrix</span><span class="p">,</span>
    <span class="n">path_node_array</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
    <span class="n">var_names</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
    <span class="n">fig_ax</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
    <span class="n">figsize</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
    <span class="n">save_name</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
    <span class="n">link_colorbar_label</span><span class="o">=</span><span class="s2">&quot;link coeff. (edge color)&quot;</span><span class="p">,</span>
    <span class="n">node_colorbar_label</span><span class="o">=</span><span class="s2">&quot;MCE (node color)&quot;</span><span class="p">,</span>
    <span class="n">link_width</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
    <span class="n">node_pos</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
    <span class="n">arrow_linewidth</span><span class="o">=</span><span class="mf">10.0</span><span class="p">,</span>
    <span class="n">vmin_edges</span><span class="o">=-</span><span class="mi">1</span><span class="p">,</span>
    <span class="n">vmax_edges</span><span class="o">=</span><span class="mf">1.0</span><span class="p">,</span>
    <span class="n">edge_ticks</span><span class="o">=</span><span class="mf">0.4</span><span class="p">,</span>
    <span class="n">cmap_edges</span><span class="o">=</span><span class="s2">&quot;RdBu_r&quot;</span><span class="p">,</span>
    <span class="n">vmin_nodes</span><span class="o">=-</span><span class="mf">1.0</span><span class="p">,</span>
    <span class="n">vmax_nodes</span><span class="o">=</span><span class="mf">1.0</span><span class="p">,</span>
    <span class="n">node_ticks</span><span class="o">=</span><span class="mf">0.4</span><span class="p">,</span>
    <span class="n">cmap_nodes</span><span class="o">=</span><span class="s2">&quot;RdBu_r&quot;</span><span class="p">,</span>
    <span class="n">node_size</span><span class="o">=</span><span class="mf">0.3</span><span class="p">,</span>
    <span class="n">node_aspect</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
    <span class="n">arrowhead_size</span><span class="o">=</span><span class="mi">20</span><span class="p">,</span>
    <span class="n">curved_radius</span><span class="o">=</span><span class="mf">0.2</span><span class="p">,</span>
    <span class="n">label_fontsize</span><span class="o">=</span><span class="mi">10</span><span class="p">,</span>
    <span class="n">tick_label_size</span><span class="o">=</span><span class="mi">6</span><span class="p">,</span>
    <span class="n">lag_array</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
    <span class="n">alpha</span><span class="o">=</span><span class="mf">1.0</span><span class="p">,</span>
    <span class="n">node_label_size</span><span class="o">=</span><span class="mi">10</span><span class="p">,</span>
    <span class="n">link_label_fontsize</span><span class="o">=</span><span class="mi">10</span><span class="p">,</span>
    <span class="c1"># network_lower_bound=0.2,</span>
    <span class="n">standard_color_links</span><span class="o">=</span><span class="s1">&#39;black&#39;</span><span class="p">,</span>
    <span class="n">standard_color_nodes</span><span class="o">=</span><span class="s1">&#39;lightgrey&#39;</span><span class="p">,</span>
<span class="p">):</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;Creates a network plot visualizing the pathways of a mediation analysis.</span>
<span class="sd">    This is still in beta. The network is defined from non-zero entries in</span>
<span class="sd">    ``path_val_matrix``.  Nodes denote variables, straight links contemporaneous</span>
<span class="sd">    dependencies and curved arrows lagged dependencies. The node color denotes</span>
<span class="sd">    the mediated causal effect (MCE) and the link color the value at the lag</span>
<span class="sd">    with maximal link coefficient. The link label lists the lags with</span>
<span class="sd">    significant dependency in order of absolute magnitude. The network can also</span>
<span class="sd">    be plotted over a map drawn before on the same axis. Then the node positions</span>
<span class="sd">    can be supplied in appropriate axis coordinates via node_pos.</span>

<span class="sd">    Parameters</span>
<span class="sd">    ----------</span>
<span class="sd">    path_val_matrix : array_like</span>
<span class="sd">        Matrix of shape (N, N, tau_max+1) containing link weight values.</span>
<span class="sd">    path_node_array: array_like</span>
<span class="sd">        Array of shape (N,) containing node values.</span>
<span class="sd">    var_names : list, optional (default: None)</span>
<span class="sd">        List of variable names. If None, range(N) is used.</span>
<span class="sd">    fig_ax : tuple of figure and axis object, optional (default: None)</span>
<span class="sd">        Figure and axes instance. If None they are created.</span>
<span class="sd">    figsize : tuple</span>
<span class="sd">        Size of figure.</span>
<span class="sd">    save_name : str, optional (default: None)</span>
<span class="sd">        Name of figure file to save figure. If None, figure is shown in window.</span>
<span class="sd">    link_colorbar_label : str, optional (default: &#39;link coeff. (edge color)&#39;)</span>
<span class="sd">        Link colorbar label.</span>
<span class="sd">    node_colorbar_label : str, optional (default: &#39;MCE (node color)&#39;)</span>
<span class="sd">        Node colorbar label.</span>
<span class="sd">    link_width : array-like, optional (default: None)</span>
<span class="sd">        Array of val_matrix.shape specifying relative link width with maximum</span>
<span class="sd">        given by arrow_linewidth. If None, all links have same width.</span>
<span class="sd">    node_pos : dictionary, optional (default: None)</span>
<span class="sd">        Dictionary of node positions in axis coordinates of form</span>
<span class="sd">        node_pos = {&#39;x&#39;:array of shape (N,), &#39;y&#39;:array of shape(N)}. These</span>
<span class="sd">        coordinates could have been transformed before for basemap plots. You can</span>
<span class="sd">        also add a key &#39;transform&#39;:ccrs.PlateCarree() in order to plot graphs on </span>
<span class="sd">        a map using cartopy.</span>
<span class="sd">    arrow_linewidth : float, optional (default: 30)</span>
<span class="sd">        Linewidth.</span>
<span class="sd">    vmin_edges : float, optional (default: -1)</span>
<span class="sd">        Link colorbar scale lower bound.</span>
<span class="sd">    vmax_edges : float, optional (default: 1)</span>
<span class="sd">        Link colorbar scale upper bound.</span>
<span class="sd">    edge_ticks : float, optional (default: 0.4)</span>
<span class="sd">        Link tick mark interval.</span>
<span class="sd">    cmap_edges : str, optional (default: &#39;RdBu_r&#39;)</span>
<span class="sd">        Colormap for links.</span>
<span class="sd">    vmin_nodes : float, optional (default: 0)</span>
<span class="sd">        Node colorbar scale lower bound.</span>
<span class="sd">    vmax_nodes : float, optional (default: 1)</span>
<span class="sd">        Node colorbar scale upper bound.</span>
<span class="sd">    node_ticks : float, optional (default: 0.4)</span>
<span class="sd">        Node tick mark interval.</span>
<span class="sd">    cmap_nodes : str, optional (default: &#39;OrRd&#39;)</span>
<span class="sd">        Colormap for links.</span>
<span class="sd">    node_size : int, optional (default: 0.3)</span>
<span class="sd">        Node size.</span>
<span class="sd">    node_aspect : float, optional (default: None)</span>
<span class="sd">        Ratio between the heigth and width of the varible nodes.</span>
<span class="sd">    arrowhead_size : int, optional (default: 20)</span>
<span class="sd">        Size of link arrow head. Passed on to FancyArrowPatch object.</span>
<span class="sd">    curved_radius, float, optional (default: 0.2)</span>
<span class="sd">        Curvature of links. Passed on to FancyArrowPatch object.</span>
<span class="sd">    label_fontsize : int, optional (default: 10)</span>
<span class="sd">        Fontsize of colorbar labels.</span>
<span class="sd">    alpha : float, optional (default: 1.)</span>
<span class="sd">        Opacity.</span>
<span class="sd">    node_label_size : int, optional (default: 10)</span>
<span class="sd">        Fontsize of node labels.</span>
<span class="sd">    link_label_fontsize : int, optional (default: 6)</span>
<span class="sd">        Fontsize of link labels.</span>
<span class="sd">    lag_array : array, optional (default: None)</span>
<span class="sd">        Optional specification of lags overwriting np.arange(0, tau_max+1)</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="n">val_matrix</span> <span class="o">=</span> <span class="n">path_val_matrix</span>

    <span class="k">if</span> <span class="n">fig_ax</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">fig</span> <span class="o">=</span> <span class="n">pyplot</span><span class="o">.</span><span class="n">figure</span><span class="p">(</span><span class="n">figsize</span><span class="o">=</span><span class="n">figsize</span><span class="p">)</span>
        <span class="n">ax</span> <span class="o">=</span> <span class="n">fig</span><span class="o">.</span><span class="n">add_subplot</span><span class="p">(</span><span class="mi">111</span><span class="p">,</span> <span class="n">frame_on</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">fig</span><span class="p">,</span> <span class="n">ax</span> <span class="o">=</span> <span class="n">fig_ax</span>

    <span class="k">if</span> <span class="n">link_width</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> <span class="ow">not</span> <span class="n">np</span><span class="o">.</span><span class="n">all</span><span class="p">(</span><span class="n">link_width</span> <span class="o">&gt;=</span> <span class="mf">0.0</span><span class="p">):</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;link_width must be non-negative&quot;</span><span class="p">)</span>

    <span class="n">N</span><span class="p">,</span> <span class="n">N</span><span class="p">,</span> <span class="n">dummy</span> <span class="o">=</span> <span class="n">val_matrix</span><span class="o">.</span><span class="n">shape</span>
    <span class="n">tau_max</span> <span class="o">=</span> <span class="n">dummy</span> <span class="o">-</span> <span class="mi">1</span>

    <span class="k">if</span> <span class="n">np</span><span class="o">.</span><span class="n">count_nonzero</span><span class="p">(</span><span class="n">val_matrix</span><span class="p">)</span> <span class="o">==</span> <span class="n">np</span><span class="o">.</span><span class="n">count_nonzero</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">diagonal</span><span class="p">(</span><span class="n">val_matrix</span><span class="p">)):</span>
        <span class="n">diagonal</span> <span class="o">=</span> <span class="kc">True</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">diagonal</span> <span class="o">=</span> <span class="kc">False</span>

    <span class="k">if</span> <span class="n">np</span><span class="o">.</span><span class="n">count_nonzero</span><span class="p">(</span><span class="n">val_matrix</span><span class="p">)</span> <span class="o">==</span> <span class="n">val_matrix</span><span class="o">.</span><span class="n">size</span> <span class="ow">or</span> <span class="n">diagonal</span><span class="p">:</span>
        <span class="n">val_matrix</span><span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">0</span><span class="p">]</span> <span class="o">=</span> <span class="mi">1</span>
        <span class="n">no_links</span> <span class="o">=</span> <span class="kc">True</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">no_links</span> <span class="o">=</span> <span class="kc">False</span>

    <span class="k">if</span> <span class="n">var_names</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">var_names</span> <span class="o">=</span> <span class="nb">range</span><span class="p">(</span><span class="n">N</span><span class="p">)</span>

    <span class="c1"># Define graph links by absolute maximum (positive or negative like for</span>
    <span class="c1"># partial correlation)</span>
    <span class="c1"># val_matrix[np.abs(val_matrix) &lt; sig_thres] = 0.</span>
    <span class="n">graph</span> <span class="o">=</span> <span class="n">val_matrix</span> <span class="o">!=</span> <span class="mf">0.0</span>
    <span class="n">net</span> <span class="o">=</span> <span class="n">_get_absmax</span><span class="p">(</span><span class="n">val_matrix</span><span class="p">)</span>
    <span class="n">G</span> <span class="o">=</span> <span class="n">nx</span><span class="o">.</span><span class="n">DiGraph</span><span class="p">(</span><span class="n">net</span><span class="p">)</span>

    <span class="n">node_color</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="n">N</span><span class="p">)</span>
    <span class="c1"># list of all strengths for color map</span>
    <span class="n">all_strengths</span> <span class="o">=</span> <span class="p">[]</span>
    <span class="c1"># Add attributes, contemporaneous and lagged links are handled separately</span>
    <span class="k">for</span> <span class="p">(</span><span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">,</span> <span class="n">dic</span><span class="p">)</span> <span class="ow">in</span> <span class="n">G</span><span class="o">.</span><span class="n">edges</span><span class="p">(</span><span class="n">data</span><span class="o">=</span><span class="kc">True</span><span class="p">):</span>
        <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;outer_edge_attribute&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span>
        <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;no_links&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">no_links</span>
        <span class="c1"># average lagfunc for link u --&gt; v ANDOR u -- v</span>
        <span class="k">if</span> <span class="n">tau_max</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
            <span class="c1"># argmax of absolute maximum</span>
            <span class="n">argmax</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">abs</span><span class="p">(</span><span class="n">val_matrix</span><span class="p">[</span><span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">][</span><span class="mi">1</span><span class="p">:])</span><span class="o">.</span><span class="n">argmax</span><span class="p">()</span> <span class="o">+</span> <span class="mi">1</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">argmax</span> <span class="o">=</span> <span class="mi">0</span>
        <span class="k">if</span> <span class="n">u</span> <span class="o">!=</span> <span class="n">v</span><span class="p">:</span>
            <span class="c1"># For contemp links masking or finite samples can lead to different</span>
            <span class="c1"># values for u--v and v--u</span>
            <span class="c1"># Here we use the  maximum for the width and weight (=color)</span>
            <span class="c1"># of the link</span>
            <span class="c1"># Draw link if u--v OR v--u at lag 0 is nonzero</span>
            <span class="c1"># dic[&#39;inner_edge&#39;] = ((np.abs(val_matrix[u, v][0]) &gt;=</span>
            <span class="c1">#                       sig_thres[u, v][0]) or</span>
            <span class="c1">#                      (np.abs(val_matrix[v, u][0]) &gt;=</span>
            <span class="c1">#                       sig_thres[v, u][0]))</span>
            <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;inner_edge&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">graph</span><span class="p">[</span><span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">,</span> <span class="mi">0</span><span class="p">]</span> <span class="ow">or</span> <span class="n">graph</span><span class="p">[</span><span class="n">v</span><span class="p">,</span> <span class="n">u</span><span class="p">,</span> <span class="mi">0</span><span class="p">]</span>
            <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;inner_edge_alpha&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">alpha</span>
            <span class="c1"># value at argmax of average</span>
            <span class="k">if</span> <span class="n">np</span><span class="o">.</span><span class="n">abs</span><span class="p">(</span><span class="n">val_matrix</span><span class="p">[</span><span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">][</span><span class="mi">0</span><span class="p">]</span> <span class="o">-</span> <span class="n">val_matrix</span><span class="p">[</span><span class="n">v</span><span class="p">,</span> <span class="n">u</span><span class="p">][</span><span class="mi">0</span><span class="p">])</span> <span class="o">&gt;</span> <span class="mf">0.0001</span><span class="p">:</span>
                <span class="nb">print</span><span class="p">(</span>
                    <span class="s2">&quot;Contemporaneous I(</span><span class="si">%d</span><span class="s2">; </span><span class="si">%d</span><span class="s2">)=</span><span class="si">%.3f</span><span class="s2"> != I(</span><span class="si">%d</span><span class="s2">; </span><span class="si">%d</span><span class="s2">)=</span><span class="si">%.3f</span><span class="s2">&quot;</span>
                    <span class="o">%</span> <span class="p">(</span><span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">,</span> <span class="n">val_matrix</span><span class="p">[</span><span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">][</span><span class="mi">0</span><span class="p">],</span> <span class="n">v</span><span class="p">,</span> <span class="n">u</span><span class="p">,</span> <span class="n">val_matrix</span><span class="p">[</span><span class="n">v</span><span class="p">,</span> <span class="n">u</span><span class="p">][</span><span class="mi">0</span><span class="p">])</span>
                    <span class="o">+</span> <span class="s2">&quot; due to conditions, finite sample effects or &quot;</span>
                    <span class="s2">&quot;masking, here edge color = &quot;</span>
                    <span class="s2">&quot;larger (absolute) value.&quot;</span>
                <span class="p">)</span>
            <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;inner_edge_color&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">_get_absmax</span><span class="p">(</span>
                <span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">([[[</span><span class="n">val_matrix</span><span class="p">[</span><span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">][</span><span class="mi">0</span><span class="p">],</span> <span class="n">val_matrix</span><span class="p">[</span><span class="n">v</span><span class="p">,</span> <span class="n">u</span><span class="p">][</span><span class="mi">0</span><span class="p">]]]])</span>
            <span class="p">)</span><span class="o">.</span><span class="n">squeeze</span><span class="p">()</span>
            <span class="k">if</span> <span class="n">link_width</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
                <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;inner_edge_width&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">arrow_linewidth</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;inner_edge_width&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="p">(</span>
                    <span class="n">link_width</span><span class="p">[</span><span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">,</span> <span class="mi">0</span><span class="p">]</span> <span class="o">/</span> <span class="n">link_width</span><span class="o">.</span><span class="n">max</span><span class="p">()</span> <span class="o">*</span> <span class="n">arrow_linewidth</span>
                <span class="p">)</span>

            <span class="n">all_strengths</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">dic</span><span class="p">[</span><span class="s2">&quot;inner_edge_color&quot;</span><span class="p">])</span>

            <span class="k">if</span> <span class="n">tau_max</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
                <span class="c1"># True if ensemble mean at lags &gt; 0 is nonzero</span>
                <span class="c1"># dic[&#39;outer_edge&#39;] = np.any(</span>
                <span class="c1">#     np.abs(val_matrix[u, v][1:]) &gt;= sig_thres[u, v][1:])</span>
                <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;outer_edge&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">any</span><span class="p">(</span><span class="n">graph</span><span class="p">[</span><span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">,</span> <span class="mi">1</span><span class="p">:])</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;outer_edge&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="kc">False</span>
            <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;outer_edge_alpha&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">alpha</span>
            <span class="k">if</span> <span class="n">link_width</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
                <span class="c1"># fraction of nonzero values</span>
                <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;outer_edge_width&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">arrow_linewidth</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;outer_edge_width&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="p">(</span>
                    <span class="n">link_width</span><span class="p">[</span><span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">,</span> <span class="n">argmax</span><span class="p">]</span> <span class="o">/</span> <span class="n">link_width</span><span class="o">.</span><span class="n">max</span><span class="p">()</span> <span class="o">*</span> <span class="n">arrow_linewidth</span>
                <span class="p">)</span>

            <span class="c1"># value at argmax of average</span>
            <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;outer_edge_color&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">val_matrix</span><span class="p">[</span><span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">][</span><span class="n">argmax</span><span class="p">]</span>
            <span class="n">all_strengths</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">dic</span><span class="p">[</span><span class="s2">&quot;outer_edge_color&quot;</span><span class="p">])</span>

            <span class="c1"># Sorted list of significant lags (only if robust wrt</span>
            <span class="c1"># d[&#39;min_ensemble_frac&#39;])</span>
            <span class="k">if</span> <span class="n">tau_max</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
                <span class="n">lags</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">abs</span><span class="p">(</span><span class="n">val_matrix</span><span class="p">[</span><span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">][</span><span class="mi">1</span><span class="p">:])</span><span class="o">.</span><span class="n">argsort</span><span class="p">()[::</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span> <span class="o">+</span> <span class="mi">1</span>
                <span class="n">sig_lags</span> <span class="o">=</span> <span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">where</span><span class="p">(</span><span class="n">graph</span><span class="p">[</span><span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">,</span> <span class="mi">1</span><span class="p">:])[</span><span class="mi">0</span><span class="p">]</span> <span class="o">+</span> <span class="mi">1</span><span class="p">)</span><span class="o">.</span><span class="n">tolist</span><span class="p">()</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">lags</span><span class="p">,</span> <span class="n">sig_lags</span> <span class="o">=</span> <span class="p">[],</span> <span class="p">[]</span>
            <span class="k">if</span> <span class="n">lag_array</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;label&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="nb">str</span><span class="p">([</span><span class="n">lag_array</span><span class="p">[</span><span class="n">l</span><span class="p">]</span> <span class="k">for</span> <span class="n">l</span> <span class="ow">in</span> <span class="n">lags</span> <span class="k">if</span> <span class="n">l</span> <span class="ow">in</span> <span class="n">sig_lags</span><span class="p">])[</span><span class="mi">1</span><span class="p">:</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">replace</span><span class="p">(</span><span class="s2">&quot; &quot;</span><span class="p">,</span> <span class="s2">&quot;&quot;</span><span class="p">)</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;label&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="nb">str</span><span class="p">([</span><span class="n">l</span> <span class="k">for</span> <span class="n">l</span> <span class="ow">in</span> <span class="n">lags</span> <span class="k">if</span> <span class="n">l</span> <span class="ow">in</span> <span class="n">sig_lags</span><span class="p">])[</span><span class="mi">1</span><span class="p">:</span><span class="o">-</span><span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">replace</span><span class="p">(</span><span class="s2">&quot; &quot;</span><span class="p">,</span> <span class="s2">&quot;&quot;</span><span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="c1"># Node color is max of average autodependency</span>
            <span class="n">node_color</span><span class="p">[</span><span class="n">u</span><span class="p">]</span> <span class="o">=</span> <span class="n">val_matrix</span><span class="p">[</span><span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">][</span><span class="n">argmax</span><span class="p">]</span>

        <span class="c1"># dic[&#39;outer_edge_edge&#39;] = False</span>
        <span class="c1"># dic[&#39;outer_edge_edgecolor&#39;] = None</span>
        <span class="c1"># dic[&#39;inner_edge_edge&#39;] = False</span>
        <span class="c1"># dic[&#39;inner_edge_edgecolor&#39;] = None</span>

    <span class="n">node_color</span> <span class="o">=</span> <span class="n">path_node_array</span>
    <span class="c1"># print node_color</span>
    <span class="c1"># If no links are present, set value to zero</span>
    <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">all_strengths</span><span class="p">)</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
        <span class="n">all_strengths</span> <span class="o">=</span> <span class="p">[</span><span class="mf">0.0</span><span class="p">]</span>

    <span class="k">if</span> <span class="n">node_pos</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">pos</span> <span class="o">=</span> <span class="n">nx</span><span class="o">.</span><span class="n">circular_layout</span><span class="p">(</span><span class="n">deepcopy</span><span class="p">(</span><span class="n">G</span><span class="p">))</span>
    <span class="c1">#            pos = nx.spring_layout(deepcopy(G))</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">pos</span> <span class="o">=</span> <span class="p">{}</span>
        <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">N</span><span class="p">):</span>
            <span class="n">pos</span><span class="p">[</span><span class="n">i</span><span class="p">]</span> <span class="o">=</span> <span class="p">(</span><span class="n">node_pos</span><span class="p">[</span><span class="s2">&quot;x&quot;</span><span class="p">][</span><span class="n">i</span><span class="p">],</span> <span class="n">node_pos</span><span class="p">[</span><span class="s2">&quot;y&quot;</span><span class="p">][</span><span class="n">i</span><span class="p">])</span>

    <span class="k">if</span> <span class="n">node_pos</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> <span class="s1">&#39;transform&#39;</span> <span class="ow">in</span> <span class="n">node_pos</span><span class="p">:</span> 
        <span class="n">transform</span> <span class="o">=</span> <span class="n">node_pos</span><span class="p">[</span><span class="s1">&#39;transform&#39;</span><span class="p">]</span>
    <span class="k">else</span><span class="p">:</span> <span class="n">transform</span> <span class="o">=</span> <span class="n">ax</span><span class="o">.</span><span class="n">transData</span>

    <span class="n">node_rings</span> <span class="o">=</span> <span class="p">{</span>
        <span class="mi">0</span><span class="p">:</span> <span class="p">{</span>
            <span class="s2">&quot;sizes&quot;</span><span class="p">:</span> <span class="kc">None</span><span class="p">,</span>
            <span class="s2">&quot;color_array&quot;</span><span class="p">:</span> <span class="n">node_color</span><span class="p">,</span>
            <span class="s2">&quot;cmap&quot;</span><span class="p">:</span> <span class="n">cmap_nodes</span><span class="p">,</span>
            <span class="s2">&quot;vmin&quot;</span><span class="p">:</span> <span class="n">vmin_nodes</span><span class="p">,</span>
            <span class="s2">&quot;vmax&quot;</span><span class="p">:</span> <span class="n">vmax_nodes</span><span class="p">,</span>
            <span class="s2">&quot;ticks&quot;</span><span class="p">:</span> <span class="n">node_ticks</span><span class="p">,</span>
            <span class="s2">&quot;label&quot;</span><span class="p">:</span> <span class="n">node_colorbar_label</span><span class="p">,</span>
            <span class="s2">&quot;colorbar&quot;</span><span class="p">:</span> <span class="kc">True</span><span class="p">,</span>
        <span class="p">}</span>
    <span class="p">}</span>

    <span class="n">_draw_network_with_curved_edges</span><span class="p">(</span>
        <span class="n">fig</span><span class="o">=</span><span class="n">fig</span><span class="p">,</span>
        <span class="n">ax</span><span class="o">=</span><span class="n">ax</span><span class="p">,</span>
        <span class="n">G</span><span class="o">=</span><span class="n">deepcopy</span><span class="p">(</span><span class="n">G</span><span class="p">),</span>
        <span class="n">pos</span><span class="o">=</span><span class="n">pos</span><span class="p">,</span>
        <span class="c1"># dictionary of rings: {0:{&#39;sizes&#39;:(N,)-array, &#39;color_array&#39;:(N,)-array</span>
        <span class="c1"># or None, &#39;cmap&#39;:string,</span>
        <span class="n">node_rings</span><span class="o">=</span><span class="n">node_rings</span><span class="p">,</span>
        <span class="c1"># &#39;vmin&#39;:float or None, &#39;vmax&#39;:float or None, &#39;label&#39;:string or None}}</span>
        <span class="n">node_labels</span><span class="o">=</span><span class="n">var_names</span><span class="p">,</span>
        <span class="n">node_label_size</span><span class="o">=</span><span class="n">node_label_size</span><span class="p">,</span>
        <span class="n">node_alpha</span><span class="o">=</span><span class="n">alpha</span><span class="p">,</span>
        <span class="n">standard_size</span><span class="o">=</span><span class="n">node_size</span><span class="p">,</span>
        <span class="n">node_aspect</span><span class="o">=</span><span class="n">node_aspect</span><span class="p">,</span>
        <span class="n">standard_cmap</span><span class="o">=</span><span class="s2">&quot;OrRd&quot;</span><span class="p">,</span>
        <span class="n">standard_color_nodes</span><span class="o">=</span><span class="n">standard_color_nodes</span><span class="p">,</span>
        <span class="n">standard_color_links</span><span class="o">=</span><span class="n">standard_color_links</span><span class="p">,</span>
        <span class="n">log_sizes</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
        <span class="n">cmap_links</span><span class="o">=</span><span class="n">cmap_edges</span><span class="p">,</span>
        <span class="n">links_vmin</span><span class="o">=</span><span class="n">vmin_edges</span><span class="p">,</span>
        <span class="n">links_vmax</span><span class="o">=</span><span class="n">vmax_edges</span><span class="p">,</span>
        <span class="n">links_ticks</span><span class="o">=</span><span class="n">edge_ticks</span><span class="p">,</span>
        <span class="n">tick_label_size</span><span class="o">=</span><span class="n">tick_label_size</span><span class="p">,</span>
        <span class="c1"># cmap_links_edges=&#39;YlOrRd&#39;, links_edges_vmin=-1., links_edges_vmax=1.,</span>
        <span class="c1"># links_edges_ticks=.2, link_edge_colorbar_label=&#39;link_edge&#39;,</span>
        <span class="n">arrowhead_size</span><span class="o">=</span><span class="n">arrowhead_size</span><span class="p">,</span>
        <span class="n">curved_radius</span><span class="o">=</span><span class="n">curved_radius</span><span class="p">,</span>
        <span class="n">label_fontsize</span><span class="o">=</span><span class="n">label_fontsize</span><span class="p">,</span>
        <span class="n">link_label_fontsize</span><span class="o">=</span><span class="n">link_label_fontsize</span><span class="p">,</span>
        <span class="n">link_colorbar_label</span><span class="o">=</span><span class="n">link_colorbar_label</span><span class="p">,</span>
        <span class="c1"># network_lower_bound=network_lower_bound,</span>
        <span class="c1"># label_fraction=label_fraction,</span>
        <span class="c1"># inner_edge_style=inner_edge_style</span>
        <span class="n">transform</span><span class="o">=</span><span class="n">transform</span>
    <span class="p">)</span>

    <span class="c1"># fig.subplots_adjust(left=0.1, right=.9, bottom=.25, top=.95)</span>
    <span class="c1"># savestring = os.path.expanduser(save_name)</span>
    <span class="k">if</span> <span class="n">save_name</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">pyplot</span><span class="o">.</span><span class="n">savefig</span><span class="p">(</span><span class="n">save_name</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">pyplot</span><span class="o">.</span><span class="n">show</span><span class="p">()</span></div>


<span class="c1">#</span>
<span class="c1">#  Functions to plot time series graphs from links including ancestors</span>
<span class="c1">#</span>
<div class="viewcode-block" id="plot_tsg"><a class="viewcode-back" href="../../index.html#tigramite.plotting.plot_tsg">[docs]</a><span class="k">def</span> <span class="nf">plot_tsg</span><span class="p">(</span><span class="n">links</span><span class="p">,</span> <span class="n">X</span><span class="p">,</span> <span class="n">Y</span><span class="p">,</span> <span class="n">Z</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">anc_x</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">anc_y</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">anc_xy</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;Plots TSG that is input in format (N*max_lag, N*max_lag).</span>
<span class="sd">    Compared to the tigramite plotting function here links</span>
<span class="sd">    X^i_{t-tau} --&gt; X^j_t can be missing for different t&#39;. Helpful to</span>
<span class="sd">    visualize the conditioned TSG.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="nf">varlag2node</span><span class="p">(</span><span class="n">var</span><span class="p">,</span> <span class="n">lag</span><span class="p">):</span>
<span class="w">        </span><span class="sd">&quot;&quot;&quot;Translate from (var, lag) notation to node in TSG.</span>
<span class="sd">        lag must be &lt;= 0.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="k">return</span> <span class="n">var</span> <span class="o">*</span> <span class="n">max_lag</span> <span class="o">+</span> <span class="n">lag</span>

    <span class="k">def</span> <span class="nf">node2varlag</span><span class="p">(</span><span class="n">node</span><span class="p">):</span>
<span class="w">        </span><span class="sd">&quot;&quot;&quot;Translate from node in TSG to (var, -tau) notation.</span>
<span class="sd">        Here tau is &lt;= 0.</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">var</span> <span class="o">=</span> <span class="n">node</span> <span class="o">//</span> <span class="n">max_lag</span>
        <span class="n">tau</span> <span class="o">=</span> <span class="n">node</span> <span class="o">%</span> <span class="p">(</span><span class="n">max_lag</span><span class="p">)</span> <span class="o">-</span> <span class="p">(</span><span class="n">max_lag</span> <span class="o">-</span> <span class="mi">1</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">var</span><span class="p">,</span> <span class="n">tau</span>

    <span class="k">def</span> <span class="nf">_get_minmax_lag</span><span class="p">(</span><span class="n">links</span><span class="p">):</span>
<span class="w">        </span><span class="sd">&quot;&quot;&quot;Helper function to retrieve tau_min and tau_max from links</span>
<span class="sd">        &quot;&quot;&quot;</span>

        <span class="n">N</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="n">links</span><span class="p">)</span>

        <span class="c1"># Get maximum time lag</span>
        <span class="n">min_lag</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">inf</span>
        <span class="n">max_lag</span> <span class="o">=</span> <span class="mi">0</span>
        <span class="k">for</span> <span class="n">j</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">N</span><span class="p">):</span>
            <span class="k">for</span> <span class="n">link_props</span> <span class="ow">in</span> <span class="n">links</span><span class="p">[</span><span class="n">j</span><span class="p">]:</span>
                <span class="n">var</span><span class="p">,</span> <span class="n">lag</span> <span class="o">=</span> <span class="n">link_props</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
                <span class="n">coeff</span> <span class="o">=</span> <span class="n">link_props</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span>
                <span class="c1"># func = link_props[2]</span>
                <span class="k">if</span> <span class="n">coeff</span> <span class="o">!=</span> <span class="mf">0.</span><span class="p">:</span>
                    <span class="n">min_lag</span> <span class="o">=</span> <span class="nb">min</span><span class="p">(</span><span class="n">min_lag</span><span class="p">,</span> <span class="nb">abs</span><span class="p">(</span><span class="n">lag</span><span class="p">))</span>
                    <span class="n">max_lag</span> <span class="o">=</span> <span class="nb">max</span><span class="p">(</span><span class="n">max_lag</span><span class="p">,</span> <span class="nb">abs</span><span class="p">(</span><span class="n">lag</span><span class="p">))</span>
        <span class="k">return</span> <span class="n">min_lag</span><span class="p">,</span> <span class="n">max_lag</span>

    <span class="k">def</span> <span class="nf">_links_to_tsg</span><span class="p">(</span><span class="n">link_coeffs</span><span class="p">,</span> <span class="n">max_lag</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
<span class="w">        </span><span class="sd">&quot;&quot;&quot;Transform link_coeffs to time series graph.</span>
<span class="sd">        TSG is of shape (N*max_lag, N*max_lag).</span>
<span class="sd">        &quot;&quot;&quot;</span>
        <span class="n">N</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="n">link_coeffs</span><span class="p">)</span>

        <span class="c1"># Get maximum lag</span>
        <span class="n">min_lag_links</span><span class="p">,</span> <span class="n">max_lag_links</span> <span class="o">=</span> <span class="n">_get_minmax_lag</span><span class="p">(</span><span class="n">link_coeffs</span><span class="p">)</span>

        <span class="c1"># max_lag of TSG is max lag in links + 1 for the zero lag.</span>
        <span class="k">if</span> <span class="n">max_lag</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">max_lag</span> <span class="o">=</span> <span class="n">max_lag_links</span> <span class="o">+</span> <span class="mi">1</span>

        <span class="n">tsg</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="n">N</span> <span class="o">*</span> <span class="n">max_lag</span><span class="p">,</span> <span class="n">N</span> <span class="o">*</span> <span class="n">max_lag</span><span class="p">))</span>

        <span class="k">for</span> <span class="n">j</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">N</span><span class="p">):</span>
            <span class="k">for</span> <span class="n">link_props</span> <span class="ow">in</span> <span class="n">link_coeffs</span><span class="p">[</span><span class="n">j</span><span class="p">]:</span>
                <span class="n">i</span><span class="p">,</span> <span class="n">lag</span> <span class="o">=</span> <span class="n">link_props</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
                <span class="n">tau</span> <span class="o">=</span> <span class="nb">abs</span><span class="p">(</span><span class="n">lag</span><span class="p">)</span>
                <span class="n">coeff</span> <span class="o">=</span> <span class="n">link_props</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span>
                <span class="c1"># func = link_props[2]</span>
                <span class="k">if</span> <span class="n">coeff</span> <span class="o">!=</span> <span class="mf">0.0</span><span class="p">:</span>
                    <span class="k">for</span> <span class="n">t</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">max_lag</span><span class="p">):</span>
                        <span class="k">if</span> <span class="p">(</span>
                            <span class="mi">0</span> <span class="o">&lt;=</span> <span class="n">varlag2node</span><span class="p">(</span><span class="n">i</span><span class="p">,</span> <span class="n">t</span> <span class="o">-</span> <span class="n">tau</span><span class="p">)</span>
                            <span class="ow">and</span> <span class="n">varlag2node</span><span class="p">(</span><span class="n">i</span><span class="p">,</span> <span class="n">t</span> <span class="o">-</span> <span class="n">tau</span><span class="p">)</span> <span class="o">%</span> <span class="n">max_lag</span>
                            <span class="o">&lt;=</span> <span class="n">varlag2node</span><span class="p">(</span><span class="n">j</span><span class="p">,</span> <span class="n">t</span><span class="p">)</span> <span class="o">%</span> <span class="n">max_lag</span>
                        <span class="p">):</span>
                            <span class="n">tsg</span><span class="p">[</span><span class="n">varlag2node</span><span class="p">(</span><span class="n">i</span><span class="p">,</span> <span class="n">t</span> <span class="o">-</span> <span class="n">tau</span><span class="p">),</span> <span class="n">varlag2node</span><span class="p">(</span><span class="n">j</span><span class="p">,</span> <span class="n">t</span><span class="p">)]</span> <span class="o">=</span> <span class="mf">1.0</span>

        <span class="k">return</span> <span class="n">tsg</span>

    <span class="n">color_list</span> <span class="o">=</span> <span class="p">[</span><span class="s2">&quot;lightgrey&quot;</span><span class="p">,</span> <span class="s2">&quot;grey&quot;</span><span class="p">,</span> <span class="s2">&quot;black&quot;</span><span class="p">,</span> <span class="s2">&quot;red&quot;</span><span class="p">,</span> <span class="s2">&quot;blue&quot;</span><span class="p">,</span> <span class="s2">&quot;orange&quot;</span><span class="p">]</span>
    <span class="n">listcmap</span> <span class="o">=</span> <span class="n">ListedColormap</span><span class="p">(</span><span class="n">color_list</span><span class="p">)</span>

    <span class="n">N</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="n">links</span><span class="p">)</span>

    <span class="n">min_lag_links</span><span class="p">,</span> <span class="n">max_lag_links</span> <span class="o">=</span> <span class="n">_get_minmax_lag</span><span class="p">(</span><span class="n">links</span><span class="p">)</span>
    <span class="n">max_lag</span> <span class="o">=</span> <span class="n">max_lag_links</span>

    <span class="k">for</span> <span class="n">anc</span> <span class="ow">in</span> <span class="n">X</span> <span class="o">+</span> <span class="n">Y</span><span class="p">:</span>
        <span class="n">max_lag</span> <span class="o">=</span> <span class="nb">max</span><span class="p">(</span><span class="n">max_lag</span><span class="p">,</span> <span class="nb">abs</span><span class="p">(</span><span class="n">anc</span><span class="p">[</span><span class="mi">1</span><span class="p">]))</span>
    <span class="k">for</span> <span class="n">anc</span> <span class="ow">in</span> <span class="n">Y</span><span class="p">:</span>
        <span class="n">max_lag</span> <span class="o">=</span> <span class="nb">max</span><span class="p">(</span><span class="n">max_lag</span><span class="p">,</span> <span class="nb">abs</span><span class="p">(</span><span class="n">anc</span><span class="p">[</span><span class="mi">1</span><span class="p">]))</span>
    <span class="k">if</span> <span class="n">Z</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="k">for</span> <span class="n">anc</span> <span class="ow">in</span> <span class="n">Z</span><span class="p">:</span>
            <span class="n">max_lag</span> <span class="o">=</span> <span class="nb">max</span><span class="p">(</span><span class="n">max_lag</span><span class="p">,</span> <span class="nb">abs</span><span class="p">(</span><span class="n">anc</span><span class="p">[</span><span class="mi">1</span><span class="p">]))</span>

    <span class="k">if</span> <span class="n">anc_x</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="k">for</span> <span class="n">anc</span> <span class="ow">in</span> <span class="n">anc_x</span><span class="p">:</span>
            <span class="n">max_lag</span> <span class="o">=</span> <span class="nb">max</span><span class="p">(</span><span class="n">max_lag</span><span class="p">,</span> <span class="nb">abs</span><span class="p">(</span><span class="n">anc</span><span class="p">[</span><span class="mi">1</span><span class="p">]))</span>
    <span class="k">if</span> <span class="n">anc_y</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="k">for</span> <span class="n">anc</span> <span class="ow">in</span> <span class="n">anc_y</span><span class="p">:</span>
            <span class="n">max_lag</span> <span class="o">=</span> <span class="nb">max</span><span class="p">(</span><span class="n">max_lag</span><span class="p">,</span> <span class="nb">abs</span><span class="p">(</span><span class="n">anc</span><span class="p">[</span><span class="mi">1</span><span class="p">]))</span>
    <span class="k">if</span> <span class="n">anc_xy</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="k">for</span> <span class="n">anc</span> <span class="ow">in</span> <span class="n">anc_xy</span><span class="p">:</span>
            <span class="n">max_lag</span> <span class="o">=</span> <span class="nb">max</span><span class="p">(</span><span class="n">max_lag</span><span class="p">,</span> <span class="nb">abs</span><span class="p">(</span><span class="n">anc</span><span class="p">[</span><span class="mi">1</span><span class="p">]))</span>

    <span class="n">max_lag</span> <span class="o">=</span> <span class="n">max_lag</span> <span class="o">+</span> <span class="mi">1</span>

    <span class="n">tsg</span> <span class="o">=</span> <span class="n">_links_to_tsg</span><span class="p">(</span><span class="n">links</span><span class="p">,</span> <span class="n">max_lag</span><span class="o">=</span><span class="n">max_lag</span><span class="p">)</span>

    <span class="n">G</span> <span class="o">=</span> <span class="n">nx</span><span class="o">.</span><span class="n">DiGraph</span><span class="p">(</span><span class="n">tsg</span><span class="p">)</span>

    <span class="n">figsize</span> <span class="o">=</span> <span class="p">(</span><span class="mi">3</span><span class="p">,</span> <span class="mi">3</span><span class="p">)</span>
    <span class="n">link_colorbar_label</span> <span class="o">=</span> <span class="s2">&quot;MCI&quot;</span>
    <span class="n">arrow_linewidth</span> <span class="o">=</span> <span class="mf">8.0</span>
    <span class="n">vmin_edges</span> <span class="o">=</span> <span class="o">-</span><span class="mi">1</span>
    <span class="n">vmax_edges</span> <span class="o">=</span> <span class="mf">1.0</span>
    <span class="n">edge_ticks</span> <span class="o">=</span> <span class="mf">0.4</span>
    <span class="n">cmap_edges</span> <span class="o">=</span> <span class="s2">&quot;RdBu_r&quot;</span>
    <span class="n">order</span> <span class="o">=</span> <span class="kc">None</span>
    <span class="n">node_size</span> <span class="o">=</span> <span class="mf">.1</span>
    <span class="n">arrowhead_size</span> <span class="o">=</span> <span class="mi">20</span>
    <span class="n">curved_radius</span> <span class="o">=</span> <span class="mf">0.2</span>
    <span class="n">label_fontsize</span> <span class="o">=</span> <span class="mi">10</span>
    <span class="n">alpha</span> <span class="o">=</span> <span class="mf">1.0</span>
    <span class="n">node_label_size</span> <span class="o">=</span> <span class="mi">10</span>
    <span class="c1"># label_space_left = 0.1</span>
    <span class="c1"># label_space_top = 0.0</span>
    <span class="c1"># network_lower_bound = 0.2</span>
    <span class="n">inner_edge_style</span> <span class="o">=</span> <span class="s2">&quot;dashed&quot;</span>

    <span class="n">node_color</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">ones</span><span class="p">(</span><span class="n">N</span> <span class="o">*</span> <span class="n">max_lag</span><span class="p">)</span>  <span class="c1"># , dtype = &#39;object&#39;)</span>
    <span class="n">node_color</span><span class="p">[:]</span> <span class="o">=</span> <span class="mi">0</span>

    <span class="k">if</span> <span class="n">anc_x</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="k">for</span> <span class="n">n</span> <span class="ow">in</span> <span class="p">[</span><span class="n">varlag2node</span><span class="p">(</span><span class="n">itau</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">max_lag</span> <span class="o">-</span> <span class="mi">1</span> <span class="o">+</span> <span class="n">itau</span><span class="p">[</span><span class="mi">1</span><span class="p">])</span> <span class="k">for</span> <span class="n">itau</span> <span class="ow">in</span> <span class="n">anc_x</span><span class="p">]:</span>
            <span class="n">node_color</span><span class="p">[</span><span class="n">n</span><span class="p">]</span> <span class="o">=</span> <span class="mi">3</span>
    <span class="k">if</span> <span class="n">anc_y</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="k">for</span> <span class="n">n</span> <span class="ow">in</span> <span class="p">[</span><span class="n">varlag2node</span><span class="p">(</span><span class="n">itau</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">max_lag</span> <span class="o">-</span> <span class="mi">1</span> <span class="o">+</span> <span class="n">itau</span><span class="p">[</span><span class="mi">1</span><span class="p">])</span> <span class="k">for</span> <span class="n">itau</span> <span class="ow">in</span> <span class="n">anc_y</span><span class="p">]:</span>
            <span class="n">node_color</span><span class="p">[</span><span class="n">n</span><span class="p">]</span> <span class="o">=</span> <span class="mi">4</span>
    <span class="k">if</span> <span class="n">anc_xy</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="k">for</span> <span class="n">n</span> <span class="ow">in</span> <span class="p">[</span><span class="n">varlag2node</span><span class="p">(</span><span class="n">itau</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">max_lag</span> <span class="o">-</span> <span class="mi">1</span> <span class="o">+</span> <span class="n">itau</span><span class="p">[</span><span class="mi">1</span><span class="p">])</span> <span class="k">for</span> <span class="n">itau</span> <span class="ow">in</span> <span class="n">anc_xy</span><span class="p">]:</span>
            <span class="n">node_color</span><span class="p">[</span><span class="n">n</span><span class="p">]</span> <span class="o">=</span> <span class="mi">5</span>

    <span class="k">for</span> <span class="n">x</span> <span class="ow">in</span> <span class="n">X</span><span class="p">:</span>
        <span class="n">node_color</span><span class="p">[</span><span class="n">varlag2node</span><span class="p">(</span><span class="n">x</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">max_lag</span> <span class="o">-</span> <span class="mi">1</span> <span class="o">+</span> <span class="n">x</span><span class="p">[</span><span class="mi">1</span><span class="p">])]</span> <span class="o">=</span> <span class="mi">2</span>
    <span class="k">for</span> <span class="n">y</span> <span class="ow">in</span> <span class="n">Y</span><span class="p">:</span>
        <span class="n">node_color</span><span class="p">[</span><span class="n">varlag2node</span><span class="p">(</span><span class="n">y</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">max_lag</span> <span class="o">-</span> <span class="mi">1</span> <span class="o">+</span> <span class="n">y</span><span class="p">[</span><span class="mi">1</span><span class="p">])]</span> <span class="o">=</span> <span class="mi">2</span>
    <span class="k">if</span> <span class="n">Z</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="k">for</span> <span class="n">z</span> <span class="ow">in</span> <span class="n">Z</span><span class="p">:</span>
            <span class="n">node_color</span><span class="p">[</span><span class="n">varlag2node</span><span class="p">(</span><span class="n">z</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">max_lag</span> <span class="o">-</span> <span class="mi">1</span> <span class="o">+</span> <span class="n">z</span><span class="p">[</span><span class="mi">1</span><span class="p">])]</span> <span class="o">=</span> <span class="mi">1</span>

    <span class="n">fig</span> <span class="o">=</span> <span class="n">pyplot</span><span class="o">.</span><span class="n">figure</span><span class="p">(</span><span class="n">figsize</span><span class="o">=</span><span class="n">figsize</span><span class="p">)</span>
    <span class="n">ax</span> <span class="o">=</span> <span class="n">fig</span><span class="o">.</span><span class="n">add_subplot</span><span class="p">(</span><span class="mi">111</span><span class="p">,</span> <span class="n">frame_on</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span>
    <span class="n">var_names</span> <span class="o">=</span> <span class="nb">range</span><span class="p">(</span><span class="n">N</span><span class="p">)</span>
    <span class="n">order</span> <span class="o">=</span> <span class="nb">range</span><span class="p">(</span><span class="n">N</span><span class="p">)</span>

    <span class="c1"># list of all strengths for color map</span>
    <span class="n">all_strengths</span> <span class="o">=</span> <span class="p">[]</span>
    <span class="c1"># Add attributes, contemporaneous and lagged links are handled separately</span>
    <span class="k">for</span> <span class="p">(</span><span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">,</span> <span class="n">dic</span><span class="p">)</span> <span class="ow">in</span> <span class="n">G</span><span class="o">.</span><span class="n">edges</span><span class="p">(</span><span class="n">data</span><span class="o">=</span><span class="kc">True</span><span class="p">):</span>
        <span class="k">if</span> <span class="n">u</span> <span class="o">!=</span> <span class="n">v</span><span class="p">:</span>
            <span class="k">if</span> <span class="n">tsg</span><span class="p">[</span><span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">]</span> <span class="ow">and</span> <span class="n">tsg</span><span class="p">[</span><span class="n">v</span><span class="p">,</span> <span class="n">u</span><span class="p">]:</span>
                <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;inner_edge&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="kc">True</span>
                <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;outer_edge&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="kc">False</span>
            <span class="k">else</span><span class="p">:</span>
                <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;inner_edge&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="kc">False</span>
                <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;outer_edge&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="kc">True</span>

            <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;inner_edge_alpha&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">alpha</span>
            <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;inner_edge_color&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">tsg</span><span class="p">[</span><span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">]</span>

            <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;inner_edge_width&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">arrow_linewidth</span>
            <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;inner_edge_attribute&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;outer_edge_attribute&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span>

            <span class="n">all_strengths</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">dic</span><span class="p">[</span><span class="s2">&quot;inner_edge_color&quot;</span><span class="p">])</span>
            <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;outer_edge_alpha&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">alpha</span>
            <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;outer_edge_width&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;inner_edge_width&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">arrow_linewidth</span>

            <span class="c1"># value at argmax of average</span>
            <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;outer_edge_color&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">tsg</span><span class="p">[</span><span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">]</span>

            <span class="n">all_strengths</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">dic</span><span class="p">[</span><span class="s2">&quot;outer_edge_color&quot;</span><span class="p">])</span>
            <span class="n">dic</span><span class="p">[</span><span class="s2">&quot;label&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span>

    <span class="c1"># If no links are present, set value to zero</span>
    <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">all_strengths</span><span class="p">)</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
        <span class="n">all_strengths</span> <span class="o">=</span> <span class="p">[</span><span class="mf">0.0</span><span class="p">]</span>

    <span class="n">posarray</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="n">N</span> <span class="o">*</span> <span class="n">max_lag</span><span class="p">,</span> <span class="mi">2</span><span class="p">))</span>
    <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">N</span> <span class="o">*</span> <span class="n">max_lag</span><span class="p">):</span>
        <span class="n">posarray</span><span class="p">[</span><span class="n">i</span><span class="p">]</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">([(</span><span class="n">i</span> <span class="o">%</span> <span class="n">max_lag</span><span class="p">),</span> <span class="p">(</span><span class="mf">1.0</span> <span class="o">-</span> <span class="n">i</span> <span class="o">//</span> <span class="n">max_lag</span><span class="p">)])</span>

    <span class="n">pos_tmp</span> <span class="o">=</span> <span class="p">{}</span>
    <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">N</span> <span class="o">*</span> <span class="n">max_lag</span><span class="p">):</span>
        <span class="n">pos_tmp</span><span class="p">[</span><span class="n">i</span><span class="p">]</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">(</span>
            <span class="p">[</span>
                <span class="p">((</span><span class="n">i</span> <span class="o">%</span> <span class="n">max_lag</span><span class="p">)</span> <span class="o">-</span> <span class="n">posarray</span><span class="o">.</span><span class="n">min</span><span class="p">(</span><span class="n">axis</span><span class="o">=</span><span class="mi">0</span><span class="p">)[</span><span class="mi">0</span><span class="p">])</span>
                <span class="o">/</span> <span class="p">(</span><span class="n">posarray</span><span class="o">.</span><span class="n">max</span><span class="p">(</span><span class="n">axis</span><span class="o">=</span><span class="mi">0</span><span class="p">)[</span><span class="mi">0</span><span class="p">]</span> <span class="o">-</span> <span class="n">posarray</span><span class="o">.</span><span class="n">min</span><span class="p">(</span><span class="n">axis</span><span class="o">=</span><span class="mi">0</span><span class="p">)[</span><span class="mi">0</span><span class="p">]),</span>
                <span class="p">((</span><span class="mf">1.0</span> <span class="o">-</span> <span class="n">i</span> <span class="o">//</span> <span class="n">max_lag</span><span class="p">)</span> <span class="o">-</span> <span class="n">posarray</span><span class="o">.</span><span class="n">min</span><span class="p">(</span><span class="n">axis</span><span class="o">=</span><span class="mi">0</span><span class="p">)[</span><span class="mi">1</span><span class="p">])</span>
                <span class="o">/</span> <span class="p">(</span><span class="n">posarray</span><span class="o">.</span><span class="n">max</span><span class="p">(</span><span class="n">axis</span><span class="o">=</span><span class="mi">0</span><span class="p">)[</span><span class="mi">1</span><span class="p">]</span> <span class="o">-</span> <span class="n">posarray</span><span class="o">.</span><span class="n">min</span><span class="p">(</span><span class="n">axis</span><span class="o">=</span><span class="mi">0</span><span class="p">)[</span><span class="mi">1</span><span class="p">]),</span>
            <span class="p">]</span>
        <span class="p">)</span>
        <span class="n">pos_tmp</span><span class="p">[</span><span class="n">i</span><span class="p">][</span><span class="n">np</span><span class="o">.</span><span class="n">isnan</span><span class="p">(</span><span class="n">pos_tmp</span><span class="p">[</span><span class="n">i</span><span class="p">])]</span> <span class="o">=</span> <span class="mf">0.0</span>

    <span class="n">pos</span> <span class="o">=</span> <span class="p">{}</span>
    <span class="k">for</span> <span class="n">n</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">N</span><span class="p">):</span>
        <span class="k">for</span> <span class="n">tau</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">max_lag</span><span class="p">):</span>
            <span class="n">pos</span><span class="p">[</span><span class="n">n</span> <span class="o">*</span> <span class="n">max_lag</span> <span class="o">+</span> <span class="n">tau</span><span class="p">]</span> <span class="o">=</span> <span class="n">pos_tmp</span><span class="p">[</span><span class="n">order</span><span class="p">[</span><span class="n">n</span><span class="p">]</span> <span class="o">*</span> <span class="n">max_lag</span> <span class="o">+</span> <span class="n">tau</span><span class="p">]</span>

    <span class="n">node_rings</span> <span class="o">=</span> <span class="p">{</span>
        <span class="mi">0</span><span class="p">:</span> <span class="p">{</span>
            <span class="s2">&quot;sizes&quot;</span><span class="p">:</span> <span class="kc">None</span><span class="p">,</span>
            <span class="s2">&quot;color_array&quot;</span><span class="p">:</span> <span class="n">node_color</span><span class="p">,</span>
            <span class="s2">&quot;label&quot;</span><span class="p">:</span> <span class="s2">&quot;&quot;</span><span class="p">,</span>
            <span class="s2">&quot;colorbar&quot;</span><span class="p">:</span> <span class="kc">False</span><span class="p">,</span>
            <span class="s2">&quot;cmap&quot;</span><span class="p">:</span> <span class="n">listcmap</span><span class="p">,</span>
            <span class="s2">&quot;vmin&quot;</span><span class="p">:</span> <span class="mi">0</span><span class="p">,</span>
            <span class="s2">&quot;vmax&quot;</span><span class="p">:</span> <span class="nb">len</span><span class="p">(</span><span class="n">color_list</span><span class="p">),</span>
        <span class="p">}</span>
    <span class="p">}</span>

    <span class="n">node_labels</span> <span class="o">=</span> <span class="p">[</span><span class="s2">&quot;&quot;</span> <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">N</span> <span class="o">*</span> <span class="n">max_lag</span><span class="p">)]</span>

    <span class="n">_draw_network_with_curved_edges</span><span class="p">(</span>
        <span class="n">fig</span><span class="o">=</span><span class="n">fig</span><span class="p">,</span>
        <span class="n">ax</span><span class="o">=</span><span class="n">ax</span><span class="p">,</span>
        <span class="n">G</span><span class="o">=</span><span class="n">deepcopy</span><span class="p">(</span><span class="n">G</span><span class="p">),</span>
        <span class="n">pos</span><span class="o">=</span><span class="n">pos</span><span class="p">,</span>
        <span class="n">node_rings</span><span class="o">=</span><span class="n">node_rings</span><span class="p">,</span>
        <span class="n">node_labels</span><span class="o">=</span><span class="n">node_labels</span><span class="p">,</span>
        <span class="n">node_label_size</span><span class="o">=</span><span class="n">node_label_size</span><span class="p">,</span>
        <span class="n">node_alpha</span><span class="o">=</span><span class="n">alpha</span><span class="p">,</span>
        <span class="n">standard_size</span><span class="o">=</span><span class="n">node_size</span><span class="p">,</span>
        <span class="n">node_aspect</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
        <span class="n">standard_cmap</span><span class="o">=</span><span class="s2">&quot;OrRd&quot;</span><span class="p">,</span>
        <span class="n">standard_color_links</span><span class="o">=</span><span class="s1">&#39;black&#39;</span><span class="p">,</span>
        <span class="n">standard_color_nodes</span><span class="o">=</span><span class="s1">&#39;lightgrey&#39;</span><span class="p">,</span>
        <span class="n">log_sizes</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
        <span class="n">cmap_links</span><span class="o">=</span><span class="n">cmap_edges</span><span class="p">,</span>
        <span class="n">links_vmin</span><span class="o">=</span><span class="n">vmin_edges</span><span class="p">,</span>
        <span class="n">links_vmax</span><span class="o">=</span><span class="n">vmax_edges</span><span class="p">,</span>
        <span class="n">links_ticks</span><span class="o">=</span><span class="n">edge_ticks</span><span class="p">,</span>
        <span class="n">arrowstyle</span><span class="o">=</span><span class="s2">&quot;simple&quot;</span><span class="p">,</span>
        <span class="n">arrowhead_size</span><span class="o">=</span><span class="n">arrowhead_size</span><span class="p">,</span>
        <span class="n">curved_radius</span><span class="o">=</span><span class="n">curved_radius</span><span class="p">,</span>
        <span class="n">label_fontsize</span><span class="o">=</span><span class="n">label_fontsize</span><span class="p">,</span>
        <span class="n">label_fraction</span><span class="o">=</span><span class="mf">0.5</span><span class="p">,</span>
        <span class="n">link_colorbar_label</span><span class="o">=</span><span class="n">link_colorbar_label</span><span class="p">,</span>
        <span class="n">inner_edge_curved</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
        <span class="c1"># network_lower_bound=network_lower_bound,</span>
        <span class="n">inner_edge_style</span><span class="o">=</span><span class="n">inner_edge_style</span><span class="p">,</span>
    <span class="p">)</span>

    <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="n">N</span><span class="p">):</span>
        <span class="n">trans</span> <span class="o">=</span> <span class="n">transforms</span><span class="o">.</span><span class="n">blended_transform_factory</span><span class="p">(</span><span class="n">ax</span><span class="o">.</span><span class="n">transAxes</span><span class="p">,</span> <span class="n">ax</span><span class="o">.</span><span class="n">transData</span><span class="p">)</span>
        <span class="n">ax</span><span class="o">.</span><span class="n">text</span><span class="p">(</span>
            <span class="mf">0.</span><span class="p">,</span>
            <span class="n">pos</span><span class="p">[</span><span class="n">order</span><span class="p">[</span><span class="n">i</span><span class="p">]</span> <span class="o">*</span> <span class="n">max_lag</span><span class="p">][</span><span class="mi">1</span><span class="p">],</span>
            <span class="s2">&quot;</span><span class="si">%s</span><span class="s2">&quot;</span> <span class="o">%</span> <span class="nb">str</span><span class="p">(</span><span class="n">var_names</span><span class="p">[</span><span class="n">order</span><span class="p">[</span><span class="n">i</span><span class="p">]]),</span>
            <span class="n">fontsize</span><span class="o">=</span><span class="n">label_fontsize</span><span class="p">,</span>
            <span class="n">horizontalalignment</span><span class="o">=</span><span class="s2">&quot;right&quot;</span><span class="p">,</span>
            <span class="n">verticalalignment</span><span class="o">=</span><span class="s2">&quot;center&quot;</span><span class="p">,</span>
            <span class="n">transform</span><span class="o">=</span><span class="n">trans</span><span class="p">,</span>
        <span class="p">)</span>

    <span class="k">for</span> <span class="n">tau</span> <span class="ow">in</span> <span class="n">np</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span><span class="n">max_lag</span> <span class="o">-</span> <span class="mi">1</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">):</span>
        <span class="n">trans</span> <span class="o">=</span> <span class="n">transforms</span><span class="o">.</span><span class="n">blended_transform_factory</span><span class="p">(</span><span class="n">ax</span><span class="o">.</span><span class="n">transData</span><span class="p">,</span> <span class="n">ax</span><span class="o">.</span><span class="n">transAxes</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">tau</span> <span class="o">==</span> <span class="n">max_lag</span> <span class="o">-</span> <span class="mi">1</span><span class="p">:</span>
            <span class="n">ax</span><span class="o">.</span><span class="n">text</span><span class="p">(</span>
                <span class="n">pos</span><span class="p">[</span><span class="n">tau</span><span class="p">][</span><span class="mi">0</span><span class="p">],</span>
                <span class="mf">1.0</span><span class="p">,</span> <span class="c1">#- label_space_top,</span>
                <span class="sa">r</span><span class="s2">&quot;$t$&quot;</span><span class="p">,</span>
                <span class="n">fontsize</span><span class="o">=</span><span class="nb">int</span><span class="p">(</span><span class="n">label_fontsize</span> <span class="o">*</span> <span class="mf">0.7</span><span class="p">),</span>
                <span class="n">horizontalalignment</span><span class="o">=</span><span class="s2">&quot;center&quot;</span><span class="p">,</span>
                <span class="n">verticalalignment</span><span class="o">=</span><span class="s2">&quot;bottom&quot;</span><span class="p">,</span>
                <span class="n">transform</span><span class="o">=</span><span class="n">trans</span><span class="p">,</span>
            <span class="p">)</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="n">ax</span><span class="o">.</span><span class="n">text</span><span class="p">(</span>
                <span class="n">pos</span><span class="p">[</span><span class="n">tau</span><span class="p">][</span><span class="mi">0</span><span class="p">],</span>
                <span class="mf">1.0</span><span class="p">,</span> <span class="c1"># - label_space_top,</span>
                <span class="sa">r</span><span class="s2">&quot;$t-</span><span class="si">%s</span><span class="s2">$&quot;</span> <span class="o">%</span> <span class="nb">str</span><span class="p">(</span><span class="n">max_lag</span> <span class="o">-</span> <span class="n">tau</span> <span class="o">-</span> <span class="mi">1</span><span class="p">),</span>
                <span class="n">fontsize</span><span class="o">=</span><span class="nb">int</span><span class="p">(</span><span class="n">label_fontsize</span> <span class="o">*</span> <span class="mf">0.7</span><span class="p">),</span>
                <span class="n">horizontalalignment</span><span class="o">=</span><span class="s2">&quot;center&quot;</span><span class="p">,</span>
                <span class="n">verticalalignment</span><span class="o">=</span><span class="s2">&quot;bottom&quot;</span><span class="p">,</span>
                <span class="n">transform</span><span class="o">=</span><span class="n">trans</span><span class="p">,</span>
            <span class="p">)</span>

    <span class="k">return</span> <span class="n">fig</span><span class="p">,</span> <span class="n">ax</span></div>

<div class="viewcode-block" id="write_csv"><a class="viewcode-back" href="../../index.html#tigramite.plotting.write_csv">[docs]</a><span class="k">def</span> <span class="nf">write_csv</span><span class="p">(</span>
    <span class="n">graph</span><span class="p">,</span>
    <span class="n">save_name</span><span class="p">,</span>
    <span class="n">val_matrix</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
    <span class="n">var_names</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
    <span class="n">link_width</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
    <span class="n">link_attribute</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
    <span class="n">digits</span><span class="o">=</span><span class="mi">5</span><span class="p">,</span>
<span class="p">):</span>
<span class="w">    </span><span class="sd">&quot;&quot;&quot;Writes all links in a graph to a csv file.</span>
<span class="sd">    </span>
<span class="sd">    Format is each link in a row as &#39;Variable i&#39;, &#39;Variable j&#39;, &#39;Time lag of i&#39;, &#39;Link type i --- j&#39;,</span>
<span class="sd">    with optional further columns for entries in [val_matrix link_attribute, link_width].</span>

<span class="sd">    Parameters</span>
<span class="sd">    ----------</span>
<span class="sd">    graph : string or bool array-like, optional (default: None)</span>
<span class="sd">        Either string matrix providing graph or bool array providing only adjacencies</span>
<span class="sd">        Must be of same shape as val_matrix. </span>
<span class="sd">    save_name : str</span>
<span class="sd">        Name of figure file to save figure. If None, figure is shown in window.</span>
<span class="sd">    val_matrix : array_like</span>
<span class="sd">        Matrix of shape (N, N, tau_max+1) containing test statistic values.</span>
<span class="sd">    var_names : list, optional (default: None)</span>
<span class="sd">        List of variable names. If None, range(N) is used.</span>
<span class="sd">    link_width : array-like, optional (default: None)</span>
<span class="sd">        Array of val_matrix.shape specifying relative link width with maximum</span>
<span class="sd">        given by arrow_linewidth. If None, all links have same width.</span>
<span class="sd">    link_attribute : array-like, optional (default: None)</span>
<span class="sd">        String array of val_matrix.shape specifying link attributes.</span>
<span class="sd">    digits : int</span>
<span class="sd">        Number of significant digits for writing link value and width.</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="n">graph</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">copy</span><span class="p">(</span><span class="n">graph</span><span class="o">.</span><span class="n">squeeze</span><span class="p">())</span>

    <span class="n">N</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="n">graph</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">val_matrix</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">val_matrix_exists</span> <span class="o">=</span> <span class="n">false</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">val_matrix_exists</span> <span class="o">=</span> <span class="kc">True</span>

    <span class="k">if</span> <span class="n">graph</span><span class="o">.</span><span class="n">ndim</span> <span class="o">==</span> <span class="mi">4</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;Time series graph of shape (N,N,tau_max+1,tau_max+1) cannot be represented by plot_graph,&quot;</span>
                         <span class="s2">&quot; use plot_time_series_graph instead.&quot;</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">graph</span><span class="o">.</span><span class="n">ndim</span> <span class="o">==</span> <span class="mi">2</span><span class="p">:</span>
        <span class="c1"># If a non-time series (N,N)-graph is given, insert a dummy dimension</span>
        <span class="n">graph</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">expand_dims</span><span class="p">(</span><span class="n">graph</span><span class="p">,</span> <span class="n">axis</span> <span class="o">=</span> <span class="mi">2</span><span class="p">)</span>

    <span class="p">(</span><span class="n">graph</span><span class="p">,</span> <span class="n">val_matrix</span><span class="p">,</span> <span class="n">link_width</span><span class="p">,</span> <span class="n">link_attribute</span><span class="p">)</span> <span class="o">=</span> <span class="n">_check_matrices</span><span class="p">(</span>
        <span class="n">graph</span><span class="p">,</span> <span class="n">val_matrix</span><span class="p">,</span> <span class="n">link_width</span><span class="p">,</span> <span class="n">link_attribute</span><span class="p">)</span>

    <span class="k">if</span> <span class="n">var_names</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">var_names</span> <span class="o">=</span> <span class="nb">range</span><span class="p">(</span><span class="n">N</span><span class="p">)</span>


    <span class="n">header</span> <span class="o">=</span> <span class="p">[</span><span class="s1">&#39;Variable i&#39;</span><span class="p">,</span> <span class="s1">&#39;Variable j&#39;</span><span class="p">,</span> <span class="s1">&#39;Time lag of i&#39;</span><span class="p">,</span> <span class="s1">&#39;Link type i --- j&#39;</span><span class="p">]</span>
    <span class="k">if</span> <span class="n">val_matrix_exists</span><span class="p">:</span>
        <span class="n">header</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;Link value&#39;</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">link_attribute</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">header</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;Link attribute&#39;</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">link_width</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
        <span class="n">header</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="s1">&#39;Link width&#39;</span><span class="p">)</span>


    <span class="k">with</span> <span class="nb">open</span><span class="p">(</span><span class="n">save_name</span><span class="p">,</span> <span class="s1">&#39;w&#39;</span><span class="p">,</span> <span class="n">encoding</span><span class="o">=</span><span class="s1">&#39;UTF8&#39;</span><span class="p">,</span> <span class="n">newline</span><span class="o">=</span><span class="s1">&#39;&#39;</span><span class="p">)</span> <span class="k">as</span> <span class="n">f</span><span class="p">:</span>
        <span class="n">writer</span> <span class="o">=</span> <span class="n">csv</span><span class="o">.</span><span class="n">writer</span><span class="p">(</span><span class="n">f</span><span class="p">)</span>

        <span class="c1"># write the header</span>
        <span class="n">writer</span><span class="o">.</span><span class="n">writerow</span><span class="p">(</span><span class="n">header</span><span class="p">)</span>

        <span class="c1"># write the link data</span>
        <span class="k">for</span> <span class="p">(</span><span class="n">i</span><span class="p">,</span> <span class="n">j</span><span class="p">,</span> <span class="n">tau</span><span class="p">)</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="o">*</span><span class="n">np</span><span class="o">.</span><span class="n">where</span><span class="p">(</span><span class="n">graph</span><span class="o">!=</span><span class="s1">&#39;&#39;</span><span class="p">)):</span>
            <span class="c1"># Only consider contemporaneous links once</span>
            <span class="k">if</span> <span class="n">tau</span> <span class="o">&gt;</span> <span class="mi">0</span> <span class="ow">or</span> <span class="n">i</span> <span class="o">&lt;=</span> <span class="n">j</span><span class="p">:</span>
                <span class="n">row</span> <span class="o">=</span> <span class="p">[</span><span class="nb">str</span><span class="p">(</span><span class="n">var_names</span><span class="p">[</span><span class="n">i</span><span class="p">]),</span> <span class="nb">str</span><span class="p">(</span><span class="n">var_names</span><span class="p">[</span><span class="n">j</span><span class="p">]),</span> <span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">tau</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">,</span> <span class="n">graph</span><span class="p">[</span><span class="n">i</span><span class="p">,</span><span class="n">j</span><span class="p">,</span><span class="n">tau</span><span class="p">]]</span>
                <span class="k">if</span> <span class="n">val_matrix_exists</span><span class="p">:</span>
                    <span class="n">row</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">val_matrix</span><span class="p">[</span><span class="n">i</span><span class="p">,</span><span class="n">j</span><span class="p">,</span><span class="n">tau</span><span class="p">]</span><span class="si">:</span><span class="s2">.</span><span class="si">{</span><span class="n">digits</span><span class="si">}}</span><span class="s2">&quot;</span><span class="p">)</span>
                <span class="k">if</span> <span class="n">link_attribute</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                    <span class="n">row</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">link_attribute</span><span class="p">[</span><span class="n">i</span><span class="p">,</span><span class="n">j</span><span class="p">,</span><span class="n">tau</span><span class="p">])</span>
                <span class="k">if</span> <span class="n">link_width</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
                    <span class="n">row</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">link_width</span><span class="p">[</span><span class="n">i</span><span class="p">,</span><span class="n">j</span><span class="p">,</span><span class="n">tau</span><span class="p">]</span><span class="si">:</span><span class="s2">.</span><span class="si">{</span><span class="n">digits</span><span class="si">}}</span><span class="s2">&quot;</span><span class="p">)</span>

                <span class="n">writer</span><span class="o">.</span><span class="n">writerow</span><span class="p">(</span><span class="n">row</span><span class="p">)</span></div>


<span class="k">if</span> <span class="vm">__name__</span> <span class="o">==</span> <span class="s2">&quot;__main__&quot;</span><span class="p">:</span>

    <span class="kn">import</span> <span class="nn">sys</span>
    <span class="n">matplotlib</span><span class="o">.</span><span class="n">rc</span><span class="p">(</span><span class="s1">&#39;xtick&#39;</span><span class="p">,</span> <span class="n">labelsize</span><span class="o">=</span><span class="mi">6</span><span class="p">)</span> 
    <span class="n">matplotlib</span><span class="o">.</span><span class="n">rc</span><span class="p">(</span><span class="s1">&#39;ytick&#39;</span><span class="p">,</span> <span class="n">labelsize</span><span class="o">=</span><span class="mi">6</span><span class="p">)</span> 

    <span class="c1"># Consider some toy data</span>
    <span class="kn">import</span> <span class="nn">tigramite</span>
    <span class="kn">import</span> <span class="nn">tigramite.toymodels.structural_causal_processes</span> <span class="k">as</span> <span class="nn">toys</span>
    <span class="kn">import</span> <span class="nn">tigramite.data_processing</span> <span class="k">as</span> <span class="nn">pp</span>
    <span class="kn">from</span> <span class="nn">tigramite.causal_effects</span> <span class="kn">import</span> <span class="n">CausalEffects</span>


    <span class="c1"># T = 1000</span>
    <span class="k">def</span> <span class="nf">lin_f</span><span class="p">(</span><span class="n">x</span><span class="p">):</span> <span class="k">return</span> <span class="n">x</span>
    <span class="c1"># auto_coeff = 0.3</span>
    <span class="c1"># coeff = 1.</span>
    <span class="c1"># links = {</span>
    <span class="c1">#         0: [((0, -1), auto_coeff, lin_f)], </span>
    <span class="c1">#         1: [((1, -1), auto_coeff, lin_f), ((0, 0), coeff, lin_f)], </span>
    <span class="c1">#         2: [((2, -1), auto_coeff, lin_f), ((1, 0), coeff, lin_f)],</span>
    <span class="c1">#         }</span>
    <span class="c1"># data, nonstat = toys.structural_causal_process(links, T=T, </span>
    <span class="c1">#                             noises=None, seed=7)</span>
    <span class="c1"># dataframe = pp.DataFrame(data, var_names=range(len(links)))</span>

    <span class="c1"># links = {</span>
    <span class="c1">#         0: [((0, -1), 1.5*auto_coeff, lin_f)], </span>
    <span class="c1">#         1: [((1, -1), 1.5*auto_coeff, lin_f), ((0, 0), 1.5*coeff, lin_f)], </span>
    <span class="c1">#         2: [((2, -1), 1.5*auto_coeff, lin_f), ((1, 0), 1.5*coeff, lin_f)],</span>
    <span class="c1">#         }</span>
    <span class="c1"># data2, nonstat = toys.structural_causal_process(links, T=T, </span>
    <span class="c1">#                             noises=None, seed=7)</span>
    <span class="c1"># dataframe2 = pp.DataFrame(data2, var_names=range(len(links)))</span>
    <span class="c1"># plot_densityplots(dataframe, name=&#39;test.pdf&#39;)</span>

    <span class="c1"># N = len(links)</span>


    <span class="c1"># parcorr = ParCorr(significance=&#39;analytic&#39;)</span>
    <span class="c1"># pcmci = PCMCI(</span>
    <span class="c1">#     dataframe=dataframe, </span>
    <span class="c1">#     cond_ind_test=parcorr,</span>
    <span class="c1">#     verbosity=1)</span>


    <span class="c1"># correlations = pcmci.get_lagged_dependencies(tau_max=20, val_only=True)[&#39;val_matrix&#39;]</span>
    <span class="c1"># lag_func_matrix = plot_lagfuncs(val_matrix=correlations, setup_args={&#39;label_space_left&#39;:0.05, </span>
    <span class="c1">#                                 &#39;x_base&#39;:5, &#39;y_base&#39;:.5})</span>
    <span class="c1"># plt.show()</span>

    
    <span class="c1"># N = len(links)</span>
    <span class="c1"># matrix = setup_density_matrix(N=N, var_names=dataframe.var_names)</span>
    <span class="c1"># matrix.add_densityplot(dataframe=dataframe, </span>
    <span class="c1">#     # selected_dataset=0, </span>
    <span class="c1">#     **{</span>
    <span class="c1">#     &#39;label&#39;:&#39;Weak&#39;,</span>
    <span class="c1">#     &#39;label_color&#39;:&#39;blue&#39;,</span>
    <span class="c1">#     &quot;snskdeplot_args&quot; : {&#39;cmap&#39;:&#39;Reds&#39;},</span>
    <span class="c1">#     }), #{&#39;cmap&#39;:&#39;Blues&#39;, &#39;alpha&#39;:0.3}})</span>
    <span class="c1"># matrix.add_densityplot(dataframe=dataframe2, selected_dataset=0, </span>
    <span class="c1">#     **{&#39;label&#39;:&#39;Strong&#39;,</span>
    <span class="c1">#     &#39;label_color&#39;:&#39;red&#39;,</span>
    <span class="c1">#     &quot;snskdeplot_args&quot; : {&#39;cmap&#39;:&#39;Blues&#39;, &#39;alpha&#39;:0.3}})</span>
    <span class="c1"># matrix.adjustfig(name=&#39;test.pdf&#39;)</span>

    <span class="c1"># matrix = setup_scatter_matrix(N=dataframe.N, </span>
    <span class="c1">#     var_names=dataframe.var_names)</span>
    <span class="c1"># matrix_lags = np.ones((3, 3)).astype(&#39;int&#39;)</span>
    <span class="c1"># matrix.add_scatterplot(dataframe=dataframe, matrix_lags=matrix_lags,</span>
    <span class="c1">#             label=&#39;ones&#39;, alpha=0.4)</span>
    <span class="c1"># matrix_lags = 2*np.ones((3, 3)).astype(&#39;int&#39;)</span>
    <span class="c1"># matrix.add_scatterplot(dataframe=dataframe, matrix_lags=matrix_lags, </span>
    <span class="c1">#     label=&#39;twos&#39;, color=&#39;red&#39;, alpha=0.4)</span>

    <span class="c1"># matrix.savefig(name=&#39;scattertest.pdf&#39;)</span>
    

    <span class="c1"># pyplot.show()</span>
    <span class="c1"># sys.exit(0)</span>


    <span class="c1"># val_matrix = np.zeros((4, 4, 3))</span>

    <span class="c1"># Complete test case</span>
    <span class="n">graph</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="mi">3</span><span class="p">,</span><span class="mi">3</span><span class="p">,</span><span class="mi">2</span><span class="p">),</span> <span class="n">dtype</span><span class="o">=</span><span class="s1">&#39;&lt;U3&#39;</span><span class="p">)</span>
    <span class="n">val_matrix</span> <span class="o">=</span> <span class="mf">0.</span><span class="o">*</span><span class="n">np</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">rand</span><span class="p">(</span><span class="o">*</span><span class="n">graph</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span>
    <span class="n">val_matrix</span><span class="p">[:,:,</span><span class="mi">0</span><span class="p">]</span> <span class="o">=</span> <span class="mf">0.2</span>
    <span class="n">graph</span><span class="p">[:]</span> <span class="o">=</span> <span class="s2">&quot;&quot;</span>
    <span class="c1"># graph[0, 1, 0] = &quot;&lt;-+&quot;</span>
    <span class="c1"># graph[1, 0, 0] = &quot;+-&gt;&quot;</span>
    <span class="n">graph</span><span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">]</span> <span class="o">=</span> <span class="s2">&quot;--&gt;&quot;</span>
    <span class="n">graph</span><span class="p">[</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">]</span> <span class="o">=</span> <span class="s2">&quot;--&gt;&quot;</span>

    <span class="n">graph</span><span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">]</span> <span class="o">=</span> <span class="s2">&quot;+-&gt;&quot;</span>
    <span class="c1"># graph[1, 0, 1] = &quot;o-o&quot;</span>

    <span class="c1"># graph[1, 2, 0] = &quot;&lt;-&gt;&quot;</span>
    <span class="c1"># graph[2, 1, 0] = &quot;&lt;-&gt;&quot;</span>

    <span class="c1"># graph[0, 2, 0] = &quot;x-x&quot;</span>
    <span class="c1"># graph[2, 0, 0] = &quot;x-x&quot;</span>
    <span class="n">nolinks</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="n">graph</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span>
    <span class="c1"># nolinks[range(4), range(4), 1] = 1</span>

    <span class="c1"># graph = graph[:2, :2, :]</span>

    <span class="c1"># fig, axes = pyplot.subplots(nrows=1, ncols=1, figsize=(6, 5))</span>


    <span class="c1"># import cartopy.crs as ccrs</span>
    <span class="n">graph</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">ones</span><span class="p">((</span><span class="mi">5</span><span class="p">,</span> <span class="mi">5</span><span class="p">,</span> <span class="mi">2</span><span class="p">),</span> <span class="n">dtype</span><span class="o">=</span><span class="s1">&#39;&lt;U3&#39;</span><span class="p">)</span>
    <span class="n">graph</span><span class="p">[:]</span> <span class="o">=</span> <span class="s2">&quot;&quot;</span>
    <span class="n">graph</span><span class="p">[</span><span class="mi">3</span><span class="p">,</span> <span class="p">:,</span> <span class="mi">1</span><span class="p">]</span> <span class="o">=</span> <span class="s1">&#39;+-&gt;&#39;</span> 

    <span class="c1"># fig = pyplot.figure(figsize=(8, 6))</span>
    <span class="c1"># fig = pyplot.figure(figsize=(10, 5))</span>
    <span class="c1"># ax = fig.add_subplot(1, 1, 1, projection=ccrs.Mollweide())</span>
    <span class="c1"># make the map global rather than have it zoom in to</span>
    <span class="c1"># the extents of any plotted data</span>
    <span class="c1"># ax.set_global()</span>
    <span class="c1"># ax.stock_img()</span>
    <span class="c1"># ax.coastlines()</span>
    <span class="c1"># # ymax = 1.</span>
    <span class="c1"># node_pos = {&#39;x&#39;:np.linspace(0, ymax, graph.shape[0]), &#39;y&#39;:np.linspace(0, ymax, graph.shape[0]),}</span>
    <span class="c1"># node_pos = {&#39;x&#39;:np.array([10,-20,80,-50,80]),</span>
    <span class="c1">#             &#39;y&#39;:np.array([-10,70,60,-40,50]), </span>
    <span class="c1">#         &#39;transform&#39;:ccrs.PlateCarree(), # t.PlateCarree()</span>
    <span class="c1">#         }</span>

    <span class="n">plot_time_series_graph</span><span class="p">(</span><span class="n">graph</span><span class="o">=</span><span class="n">graph</span><span class="p">,</span>
        <span class="c1"># fig_ax = (fig, ax),</span>
        <span class="c1"># val_matrix=val_matrix,</span>
        <span class="c1"># figsize=(5, 5),</span>
        <span class="c1"># var_names = [&#39;Var %s&#39; %i for i in range(len(graph))],</span>
        <span class="c1"># arrow_linewidth=6,</span>
        <span class="c1"># label_space_left = label_space_left,</span>
        <span class="c1"># label_space_top = label_space_top,</span>
        <span class="c1"># # network_lower_bound=network_lower_bound,</span>
        <span class="n">save_name</span><span class="o">=</span><span class="s2">&quot;tsg_test.pdf&quot;</span>
        <span class="p">)</span>
    <span class="n">pyplot</span><span class="o">.</span><span class="n">tight_layout</span><span class="p">()</span>

    <span class="c1"># network_lower_bound = 0.</span>
    <span class="c1"># show_colorbar=True</span>
    <span class="c1"># plot_graph(graph=graph,</span>
    <span class="c1">#     fig_ax = (fig, ax),</span>
    <span class="c1">#     node_pos = node_pos,</span>
    <span class="c1">#     node_size = 20,</span>
    <span class="c1">#     # val_matrix=val_matrix,</span>
    <span class="c1">#     # figsize=(5, 5),</span>
    <span class="c1">#     # var_names = [&#39;Var %s&#39; %i for i in range(len(graph))],</span>
    <span class="c1">#     # arrow_linewidth=6,</span>
    <span class="c1">#     # label_space_left = label_space_left,</span>
    <span class="c1">#     # label_space_top = label_space_top,</span>
    <span class="c1">#     # # network_lower_bound=network_lower_bound,</span>
    <span class="c1">#     save_name=&quot;tsg_test.pdf&quot;</span>
    <span class="c1">#     )</span>
    <span class="c1"># pyplot.tight_layout()</span>
    <span class="c1"># axes[0,0].scatter(np.random.rand(100), np.random.rand(100))</span>

    <span class="c1"># plot_graph(graph=graph,</span>
    <span class="c1">#     fig_ax = (fig, axes[0,0]),</span>
    <span class="c1">#     val_matrix=val_matrix,</span>
    <span class="c1">#     # figsize=(5, 5),</span>
    <span class="c1">#     var_names = [&#39;Variable %s&#39; %i for i in range(len(graph))],</span>
    <span class="c1">#     arrow_linewidth=6,</span>
    <span class="c1">#     # label_space_left = label_space_left,</span>
    <span class="c1">#     # label_space_top = label_space_top,</span>
    <span class="c1">#     # save_name=&quot;tsg_test.pdf&quot;</span>
    <span class="c1">#     )</span>
    <span class="c1"># plot_graph(graph=graph,</span>
    <span class="c1">#     fig_ax = (fig, axes[0,1]),</span>
    <span class="c1">#     val_matrix=val_matrix,</span>
    <span class="c1">#     var_names = [&#39;Var %s&#39; %i for i in range(len(graph))],</span>
    <span class="c1">#     arrow_linewidth=6,</span>
    <span class="c1">#     # label_space_left = label_space_left,</span>
    <span class="c1">#     # label_space_top = label_space_top,</span>
    <span class="c1">#     )</span>
    <span class="c1"># plot_graph(graph=graph,</span>
    <span class="c1">#     fig_ax = (fig, axes[1,0]),</span>
    <span class="c1">#     val_matrix=val_matrix,</span>
    <span class="c1">#     var_names = [&#39;Var %s&#39; %i for i in range(len(graph))],</span>
    <span class="c1">#     arrow_linewidth=6,</span>
    <span class="c1">#     # label_space_left = label_space_left,</span>
    <span class="c1">#     # label_space_top = label_space_top,</span>
    <span class="c1">#     )</span>
    <span class="c1"># plot_graph(graph=graph,</span>
    <span class="c1">#     fig_ax = (fig, axes[1,1]),</span>
    <span class="c1">#     val_matrix=val_matrix,</span>
    <span class="c1">#     var_names = [&#39;Var %s&#39; %i for i in range(len(graph))],</span>
    <span class="c1">#     arrow_linewidth=6,</span>
    <span class="c1">#     n</span>
    <span class="c1">#     # label_space_left = label_space_left,</span>
    <span class="c1">#     # label_space_top = label_space_top,</span>
    <span class="c1">#     )</span>
    <span class="c1"># # pyplot.subplots_adjust(wspace=0.3, hspace=0.2)</span>
    <span class="c1"># pyplot.tight_layout()</span>
    <span class="c1"># pyplot.savefig(&quot;test.pdf&quot;)</span>

    <span class="c1"># def lin_f(x): return x</span>

    <span class="c1"># links_coeffs = {0: [((0, -1), 0.3, lin_f)], #, ((1, -1), 0.5, lin_f)],</span>
    <span class="c1">#             1: [((1, -1), 0.3, lin_f), ((0, 0), 0.7, lin_f), ((2, -1), 0.5, lin_f)],</span>
    <span class="c1">#             2: [],</span>
    <span class="c1">#             3: [((3, -1), 0., lin_f), ((2, 0), 0.6, lin_f),]</span>
    <span class="c1">#             }</span>
    <span class="c1"># graph = CausalEffects.get_graph_from_dict(links_coeffs, tau_max=None)</span>

    <span class="c1"># val_matrix = np.random.randn(*graph.shape)</span>
    <span class="c1"># val_matrix[:,:,0] = 0.</span>
    <span class="c1"># write_csv(graph=graph,</span>
    <span class="c1">#     val_matrix=val_matrix,</span>
    <span class="c1">#     var_names=[r&#39;$X^{%d}$&#39; %i for i in range(graph.shape[0])],</span>
    <span class="c1">#     link_width=np.ones(graph.shape),</span>
    <span class="c1">#     link_attribute = np.ones(graph.shape, dtype=&#39;&lt;U10&#39;),</span>
    <span class="c1">#     save_name=&#39;test.csv&#39;)</span>

    <span class="c1"># # print(graph)</span>
    <span class="c1"># X = [(0,-1)]</span>
    <span class="c1"># Y = [(1,0)]</span>
    <span class="c1"># causal_effects = CausalEffects(graph, graph_type=&#39;stationary_dag&#39;, X=X, Y=Y, S=None, </span>
    <span class="c1">#                                hidden_variables=[(2, 0), (2, -1), (2, -2)], </span>
    <span class="c1">#                                verbosity=0)</span>

    <span class="c1"># pyplot.show()</span>
</pre></div>

          </div>
          
        </div>
      </div>
      <div class="sphinxsidebar" role="navigation" aria-label="main navigation">
        <div class="sphinxsidebarwrapper">
<h1 class="logo"><a href="../../index.html">Tigramite</a></h1>








<h3>Navigation</h3>

<div class="relations">
<h3>Related Topics</h3>
<ul>
  <li><a href="../../index.html">Documentation overview</a><ul>
  <li><a href="../index.html">Module code</a><ul>
  </ul></li>
  </ul></li>
</ul>
</div>
<div id="searchbox" style="display: none" role="search">
  <h3 id="searchlabel">Quick search</h3>
    <div class="searchformwrapper">
    <form class="search" action="../../search.html" method="get">
      <input type="text" name="q" aria-labelledby="searchlabel" autocomplete="off" autocorrect="off" autocapitalize="off" spellcheck="false"/>
      <input type="submit" value="Go" />
    </form>
    </div>
</div>
<script>document.getElementById('searchbox').style.display = "block"</script>








        </div>
      </div>
      <div class="clearer"></div>
    </div>
    <div class="footer">
      &copy;2023, Jakob Runge.
      
      |
      Powered by <a href="http://sphinx-doc.org/">Sphinx 5.0.2</a>
      &amp; <a href="https://github.com/bitprophet/alabaster">Alabaster 0.7.12</a>
      
    </div>

    

    
  </body>
</html>